download extension from rust works

ALEK_LIST_FILES.txt

@@ -1 +0,0 @@
-[RemotePath("v14/share/postgresql/extension/test_ext.control")]

ALEK_SHAREDIR.txt

@@ -1 +0,0 @@
-"v14/share/postgresql/extension"

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<GenericRem
     let remote_ext_config: serde_json::Value = serde_json::from_str(remote_ext_config)?;
     let remote_ext_bucket = match &remote_ext_config["bucket"] {
         Value::String(x) => 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,

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<Vec<RemotePath>> {
+        // 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());

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,
         }
     )