rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-01-15 09:22:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

added support for private libraries. refactored library downloading function to be more efficient by using a hashmap and to reduce code duplication

Browse Source
This commit is contained in:
Alek Westover
2023-06-27 16:56:11 -04:00
parent 5e882d8359
commit b402dfd9c7
3 changed files with 69 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
compute_tools/src/compute.rs
View File

@@ -698,12 +698,15 @@ LIMIT 100",
// TODO write a proper test for this
// Currently pytest doesn't pass cluster settings to compute_ctl
// We need to add this to pytest.
//
// libs_vec.push("test_lib1".to_string());
// info!(
// "shared_preload_libraries extra settings set to {:?}",
// libs_vec
// );
// and neon_local pass to spec
libs_vec.push("test_lib1".to_string());
libs_vec.push("private_lib1".to_string());
libs_vec.push("test_lib0".to_string());
libs_vec.push("private_lib0".to_string());
info!(
"shared_preload_libraries extra settings set to {:?}",
libs_vec
);
// download extension control files & shared_preload_libraries
@@ -714,6 +717,7 @@ LIMIT 100",
)
.await?;
info!("Libraries to download: {:?}", &libs_vec);
extension_server::get_available_libraries(
ext_remote_storage,
&self.pgbin,

96
compute_tools/src/extension_server.rs
View File

@@ -1,6 +1,7 @@
use anyhow::{self, bail, Result};
use remote_storage::*;
use serde_json::{self, Value};
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::num::{NonZeroU32, NonZeroUsize};
@@ -132,83 +133,50 @@ pub async fn get_available_libraries(
let local_libdir: PathBuf = Path::new(&get_pg_config("--pkglibdir", pgbin)).into();
let pg_version = get_pg_version(pgbin);
let remote_libdir = RemotePath::new(&Path::new(&pg_version).join("lib/")).unwrap();
// Construct a hashmap of all available libraries
// example (key, value) pair: test_lib0.so, v14/lib/test_lib0.so
let mut all_available_libraries = HashMap::new();
// 1. Download public libraries
let available_libraries = remote_storage.list_files(Some(&remote_libdir)).await?;
info!("list of library files {:?}", &available_libraries);
let remote_libdir_public = RemotePath::new(&Path::new(&pg_version).join("lib/")).unwrap();
for public_lib in remote_storage
.list_files(Some(&remote_libdir_public))
.await?
{
all_available_libraries.insert(
public_lib.object_name().expect("bad object").to_owned(),
public_lib.clone().to_owned(),
);
}
for private_prefix in private_ext_prefixes {
let remote_libdir_private =
RemotePath::new(&Path::new(&pg_version).join(private_prefix).join("lib")).unwrap();
for private_lib in remote_storage
.list_files(Some(&remote_libdir_private))
.await?
{
all_available_libraries.insert(
private_lib.object_name().expect("bad object").to_owned(),
private_lib.to_owned(),
);
}
}
// download all requested libraries
// add file extension if it isn't in the filename
for lib_name in preload_libraries {
// add file extension if it isn't in the filename
let lib_name_with_ext = if !lib_name.ends_with(".so") {
lib_name.to_owned() + ".so"
} else {
lib_name.to_string()
};
info!("looking for library {:?}", &lib_name_with_ext);
for lib in available_libraries.iter() {
info!("object_name {}", lib.object_name().unwrap());
}
let lib_path = available_libraries
.iter()
.find(|lib: &&RemotePath| lib.object_name().unwrap() == lib_name_with_ext);
match lib_path {
// TODO don't panic here,
// remember error and return it only if library is not found in any prefix
match all_available_libraries.get(&*lib_name_with_ext) {
Some(remote_path) => {
download_helper(remote_storage, remote_path, &local_libdir).await?
}
None => bail!("Shared library file {lib_name} is not found in the extension store"),
Some(lib_path) => {
download_helper(remote_storage, lib_path, &local_libdir).await?;
info!("downloaded library {:?}", &lib_path);
}
}
}
// 2. Download private libraries
for private_prefix in private_ext_prefixes {
let remote_libdir_private =
RemotePath::new(&Path::new(&pg_version).join(private_prefix).join("lib/")).unwrap();
let available_libraries_private = remote_storage
.list_files(Some(&remote_libdir_private))
.await?;
info!("list of library files {:?}", &available_libraries_private);
// download all requested libraries
// add file extension if it isn't in the filename
//
// TODO refactor this code to avoid duplication
for lib_name in preload_libraries {
let lib_name_with_ext = if !lib_name.ends_with(".so") {
lib_name.to_owned() + ".so"
} else {
lib_name.to_string()
};
info!("looking for library {:?}", &lib_name_with_ext);
for lib in available_libraries_private.iter() {
info!("object_name {}", lib.object_name().unwrap());
}
let lib_path = available_libraries_private
.iter()
.find(|lib: &&RemotePath| lib.object_name().unwrap() == lib_name_with_ext);
match lib_path {
None => bail!("Shared library file {lib_name} is not found in the extension store"),
Some(lib_path) => {
download_helper(remote_storage, lib_path, &local_libdir).await?;
info!("downloaded library {:?}", &lib_path);
}
}
}
}
Ok(())
}

45
test_runner/regress/test_download_extensions.py
View File

@@ -65,16 +65,12 @@ def prepare_mock_ext_storage(
# Upload several test_ext{i}.control files to the bucket
for i in range(NUM_EXT):
# public extensions
public_ext = BytesIO(bytes(control_file_content("public", i), "utf-8"))
public_remote_name = f"{bucket_prefix}/{PUB_EXT_ROOT}/test_ext{i}.control"
public_local_name = f"{LOCAL_EXT_ROOT}/test_ext{i}.control"
# private extensions
private_ext = BytesIO(bytes(control_file_content(str(tenant_id), i), "utf-8"))
private_remote_name = f"{bucket_prefix}/{PRIVATE_EXT_ROOT}/private_ext{i}.control"
private_local_name = f"{LOCAL_EXT_ROOT}/private_ext{i}.control"
cleanup_files += [public_local_name, private_local_name]
remote_storage_client.upload_fileobj(
@@ -97,29 +93,32 @@ def prepare_mock_ext_storage(
cleanup_files += [test_sql_local_path]
# upload some fake library files
# TODO change it to test both public and private library paths
for i in range(2):
lib_filename = f"test_lib{i}.so"
TEST_LIB_PATH = f"{PUB_LIB_ROOT}/{lib_filename}"
lib_public_remote_path = f"{bucket_prefix}/{TEST_LIB_PATH}"
lib_local_path = f"{LOCAL_LIB_ROOT}/{lib_filename}"
test_lib_file = BytesIO(
b"""
111
"""
public_library = BytesIO(bytes("\n111\n", "utf-8"))
public_remote_name = f"{bucket_prefix}/{PUB_LIB_ROOT}/test_lib{i}.so"
public_local_name = f"{LOCAL_LIB_ROOT}/test_lib{i}.so"
private_library = BytesIO(bytes("\n111\n", "utf-8"))
private_remote_name = f"{bucket_prefix}/{PRIVATE_LIB_ROOT}/private_lib{i}.so"
private_local_name = f"{LOCAL_EXT_ROOT}/private_lib{i}.so"
log.info(f"uploading library to {public_remote_name}")
log.info(f"uploading library to {private_remote_name}")
remote_storage_client.upload_fileobj(
public_library,
ext_remote_storage.bucket_name,
public_remote_name,
)
remote_storage_client.upload_fileobj(
test_lib_file,
private_library,
ext_remote_storage.bucket_name,
lib_public_remote_path,
private_remote_name,
)
log.info(f"lib_local_path: {lib_local_path}")
cleanup_files += [lib_local_path]
cleanup_files += [public_local_name, private_local_name]
return cleanup_files
#
# Generate mock extension files and upload them to the bucket.
#
# Then check that compute nodes can download them and use them
@@ -202,6 +201,16 @@ def test_remote_extensions(
log.info("LOAD test_lib0.so failed (expectedly): %s", e)
assert "file too short" in str(e)
# Try to load private library file
try:
cur.execute("LOAD 'private_lib0.so'")
except Exception as e:
# expected to fail with
# could not load library ... test_ext.so: file too short
# because test_lib0.so is not real library file
log.info("LOAD private_lib0.so failed (expectedly): %s", e)
assert "file too short" in str(e)
# Try to load existing library file without .so extension
try:
cur.execute("LOAD 'test_lib1'")