add support for public extensions. add support for library_index (libraries with names different from the archive they point to)

compute_tools/src/compute.rs

@@ -58,6 +58,7 @@ pub struct ComputeNode {
     pub ext_remote_storage: Option<GenericRemoteStorage>,
     // (key: extension name, value: path to extension archive in remote storage)
     pub ext_remote_paths: OnceLock<HashMap<String, RemotePath>>,
+    pub library_index: OnceLock<HashMap<String, String>>,
     pub already_downloaded_extensions: Mutex<HashSet<String>>,
     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 {

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<HashMap<String, RemotePath>> {
+) -> Result<(HashMap<String, RemotePath>, HashMap<String, String>)> {
     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::<Vec<&str>>();
+    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,

compute_tools/src/http/api.rs

@@ -125,13 +125,23 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
         (&Method::POST, route) if route.starts_with("/extension_server/") => {
             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);

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"}}
+{
+    "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"
+        }
+    }
+}

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"}}
+{
+    "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"
+        }
+    }
+}