From 4b83a206bf4777604df28df64ffd70f577a9b73d Mon Sep 17 00:00:00 2001 From: Alek Westover Date: Wed, 21 Jun 2023 17:22:21 -0400 Subject: [PATCH] download extension from rust works --- ALEK_LIST_FILES.txt | 1 - ALEK_SHAREDIR.txt | 1 - compute_tools/src/extension_server.rs | 34 +++++++++---------- libs/remote_storage/src/s3_bucket.rs | 1 + .../regress/test_download_extensions.py | 11 +++--- 5 files changed, 24 insertions(+), 24 deletions(-) delete mode 100644 ALEK_LIST_FILES.txt delete mode 100644 ALEK_SHAREDIR.txt diff --git a/ALEK_LIST_FILES.txt b/ALEK_LIST_FILES.txt deleted file mode 100644 index 698472fafd..0000000000 --- a/ALEK_LIST_FILES.txt +++ /dev/null @@ -1 +0,0 @@ -[RemotePath("v14/share/postgresql/extension/test_ext.control")] diff --git a/ALEK_SHAREDIR.txt b/ALEK_SHAREDIR.txt deleted file mode 100644 index 9735be8b4a..0000000000 --- a/ALEK_SHAREDIR.txt +++ /dev/null @@ -1 +0,0 @@ -"v14/share/postgresql/extension" \ No newline at end of file diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs index b749f6e927..d07002cfb7 100644 --- a/compute_tools/src/extension_server.rs +++ b/compute_tools/src/extension_server.rs @@ -1,4 +1,4 @@ -use anyhow::{self}; +use anyhow::{self, bail}; use remote_storage::*; use serde_json::{self, Value}; use std::fs::File; @@ -43,10 +43,12 @@ async fn download_helper( remote_from_path: &RemotePath, download_to_dir: &str, ) -> anyhow::Result<()> { - std::fs::write("ALEK_DOWNLOAD.txt", format!("{:?}", download_to_dir))?; - let file_name = remote_from_path.object_name().expect("it must exist"); - info!("Downloading {:?}", file_name); - info!("To location {:?}", download_to_dir); + let file_name = remote_from_path.with_base(Path::new("pg_install")); + info!( + "Downloading {:?} to location {:?}", + file_name, + file_name.clone() + ); let mut download = remote_storage.download(&remote_from_path).await?; let mut write_data_buffer = Vec::new(); download @@ -69,13 +71,9 @@ pub async fn download_extension( ext_type: ExtensionType, pgbin: &str, ) -> anyhow::Result<()> { - let from_paths = remote_storage.list_files(None).await?; - std::fs::write("ALEK_LIST_FILES.txt", format!("{:?}", from_paths))?; - let (mut local_sharedir, mut remote_sharedir) = get_pg_config("--sharedir", pgbin); local_sharedir.push_str("/extension"); remote_sharedir.push_str("/extension"); - std::fs::write("ALEK_SHAREDIR.txt", format!("{:?}", remote_sharedir))?; let (local_libdir, _) = get_pg_config("--libdir", pgbin); match ext_type { ExtensionType::Shared => { @@ -84,10 +82,6 @@ pub async fn download_extension( // because public extensions are common for all projects. let folder = RemotePath::new(Path::new(&remote_sharedir))?; let from_paths = remote_storage.list_files(Some(&folder)).await?; - std::fs::write( - "ALEK_QUEUE_DOWNLOAD.txt", - format!("{:?}", from_paths.clone()), - )?; for remote_from_path in from_paths { if remote_from_path.extension() == Some("control") { download_helper(&remote_storage, &remote_from_path, &local_sharedir).await?; @@ -120,22 +114,26 @@ pub fn init_remote_storage(remote_ext_config: &str) -> anyhow::Result x, - _ => panic!("oops"), + _ => bail!("remote_ext_config missing bucket"), }; let remote_ext_region = match &remote_ext_config["region"] { Value::String(x) => x, - _ => panic!("oops"), + _ => bail!("remote_ext_config missing region"), }; let remote_ext_endpoint = match &remote_ext_config["endpoint"] { Value::String(x) => Some(x.clone()), _ => None, }; + let remote_ext_prefix = match &remote_ext_config["prefix"] { + Value::String(x) => Some(x.clone()), + _ => None, + }; // load will not be large, so default parameters are fine let config = S3Config { - bucket_name: remote_ext_bucket.clone(), - bucket_region: remote_ext_region.clone(), - prefix_in_bucket: None, + bucket_name: remote_ext_bucket.to_string(), + bucket_region: remote_ext_region.to_string(), + prefix_in_bucket: remote_ext_prefix, endpoint: remote_ext_endpoint, concurrency_limit: NonZeroUsize::new(100).expect("100 != 0"), max_keys_per_list_response: None, diff --git a/libs/remote_storage/src/s3_bucket.rs b/libs/remote_storage/src/s3_bucket.rs index 43d818dfb9..2904e8fe95 100644 --- a/libs/remote_storage/src/s3_bucket.rs +++ b/libs/remote_storage/src/s3_bucket.rs @@ -349,6 +349,7 @@ impl RemoteStorage for S3Bucket { /// See the doc for `RemoteStorage::list_files` async fn list_files(&self, folder: Option<&RemotePath>) -> anyhow::Result> { + // TODO: if bucket prefix is empty, folder is prefixed with a "/" I think. Is this desired? let folder_name = folder .map(|p| self.relative_path_to_s3_object(p)) .or_else(|| self.prefix_in_bucket.clone()); diff --git a/test_runner/regress/test_download_extensions.py b/test_runner/regress/test_download_extensions.py index 3154abd6b2..95298a03d4 100644 --- a/test_runner/regress/test_download_extensions.py +++ b/test_runner/regress/test_download_extensions.py @@ -5,6 +5,7 @@ from fixtures.neon_fixtures import ( RemoteStorageKind, ) import json +import os def test_file_download(neon_env_builder: NeonEnvBuilder): @@ -27,23 +28,24 @@ def test_file_download(neon_env_builder: NeonEnvBuilder): env = neon_env_builder.init_start() TEST_EXT_PATH = "v14/share/postgresql/extension/test_ext.control" + BUCKET_PREFIX = "5314225671" # this is a hash of the commit number # 4. Upload test_ext.control file to the bucket # In the non-mock version this is done by CI/CD with open("test_ext.control", "rb") as data: env.remote_storage_client.upload_fileobj( - data, env.ext_remote_storage.bucket_name, TEST_EXT_PATH + data, env.ext_remote_storage.bucket_name, os.path.join(BUCKET_PREFIX, TEST_EXT_PATH) ) # 5. Download file from the bucket to correct local location # Later this will be replaced by our rust code resp = env.remote_storage_client.get_object( - Bucket=env.ext_remote_storage.bucket_name, Key=TEST_EXT_PATH + Bucket=env.ext_remote_storage.bucket_name, Key=os.path.join(BUCKET_PREFIX, TEST_EXT_PATH) ) response = resp["Body"] fname = f"pg_install/{TEST_EXT_PATH}" - with open(fname, "wb") as f: - f.write(response.read()) + # with open(fname, "wb") as f: + # f.write(response.read()) tenant, _ = env.neon_cli.create_tenant() env.neon_cli.create_timeline("test_file_download", tenant_id=tenant) @@ -53,6 +55,7 @@ def test_file_download(neon_env_builder: NeonEnvBuilder): "bucket": env.ext_remote_storage.bucket_name, "region": "us-east-1", "endpoint": env.ext_remote_storage.endpoint, + "prefix": BUCKET_PREFIX, } )