From 0b11d8e83637d72395be2c55295e59f8ff93c10a Mon Sep 17 00:00:00 2001
From: Alek Westover <alek.westover@gmail.com>
Date: Wed, 21 Jun 2023 14:01:44 -0400
Subject: [PATCH] replaced download_files function with more appropriate
 download_extensions function

---
 compute_tools/src/bin/compute_ctl.rs  | 11 +++++++----
 compute_tools/src/compute.rs          |  6 ++++--
 compute_tools/src/extension_server.rs |  7 -------
 compute_tools/src/http/api.rs         |  9 +++++++--
 4 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/compute_tools/src/bin/compute_ctl.rs b/compute_tools/src/bin/compute_ctl.rs
index e937a9ba2b..39eb6e9173 100644
--- a/compute_tools/src/bin/compute_ctl.rs
+++ b/compute_tools/src/bin/compute_ctl.rs
@@ -49,7 +49,7 @@ use compute_api::responses::ComputeStatus;
 
 use compute_tools::compute::{ComputeNode, ComputeState, ParsedSpec};
 use compute_tools::configurator::launch_configurator;
-use compute_tools::extension_server::download_file;
+use compute_tools::extension_server::{download_extension, init_remote_storage, ExtensionType};
 use compute_tools::http::api::launch_http_server;
 use compute_tools::logger::*;
 use compute_tools::monitor::launch_monitor;
@@ -66,12 +66,15 @@ fn main() -> Result<()> {
     let remote_ext_config = matches
         .get_one::<String>("remote-ext-config")
         .expect("remote-extension-config is required");
+    let remote_storage = init_remote_storage(remote_ext_config)?;
 
+    // TODO: can we give remote_storage a static lifetime, so that we don't have to copy it?
+    let copy_remote_storage = remote_storage.clone();
     let rt = Runtime::new().unwrap();
     rt.block_on(async move {
-        download_file("test_ext.control", remote_ext_config)
+        download_extension(&copy_remote_storage, ExtensionType::Shared)
             .await
-            .expect("download should work");
+            .expect("download extension should work");
     });
 
     let http_port = *matches
@@ -191,7 +194,7 @@ fn main() -> Result<()> {
         live_config_allowed,
         state: Mutex::new(new_state),
         state_changed: Condvar::new(),
-        remote_ext_config: remote_ext_config.clone(),
+        remote_storage: remote_storage,
     };
     let compute = Arc::new(compute_node);
 
diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index 1e166c45b0..3991b27196 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -16,6 +16,8 @@ use utils::lsn::Lsn;
 use compute_api::responses::{ComputeMetrics, ComputeStatus};
 use compute_api::spec::{ComputeMode, ComputeSpec};
 
+use remote_storage::GenericRemoteStorage;
+
 use crate::config;
 use crate::pg_helpers::*;
 use crate::spec::*;
@@ -45,8 +47,8 @@ pub struct ComputeNode {
     pub state: Mutex<ComputeState>,
     /// `Condvar` to allow notifying waiters about state changes.
     pub state_changed: Condvar,
-    //  S3 configuration variables:
-    pub remote_ext_config: String,
+    ///  S3 extensions configuration variables (JSON)
+    pub remote_storage: GenericRemoteStorage,
 }
 
 #[derive(Clone, Debug)]
diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs
index f1ff9034f7..f849e8c677 100644
--- a/compute_tools/src/extension_server.rs
+++ b/compute_tools/src/extension_server.rs
@@ -9,13 +9,6 @@ use std::str;
 use tokio::io::AsyncReadExt;
 use tracing::info;
 
-// TODO: get rid of this function by making s3_config part of ComputeNode
-pub async fn download_file(filename: &str, remote_ext_config: &str) -> anyhow::Result<()> {
-    let remote_storage = init_remote_storage(remote_ext_config)?;
-    download_extension(&remote_storage, ExtensionType::Shared).await?;
-    Ok(())
-}
-
 fn get_pg_config(argument: &str) -> String {
     let config_output = std::process::Command::new("pg_config")
         .arg(argument)
diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs
index f936e28e17..bb80e4272e 100644
--- a/compute_tools/src/http/api.rs
+++ b/compute_tools/src/http/api.rs
@@ -16,7 +16,7 @@ use tokio::task;
 use tracing::{error, info};
 use tracing_utils::http::OtelName;
 
-use crate::extension_server;
+use crate::extension_server::{self, ExtensionType};
 
 fn status_response_from_state(state: &ComputeState) -> ComputeStatusResponse {
     ComputeStatusResponse {
@@ -134,7 +134,12 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
                 filename
             );
 
-            match extension_server::download_file(filename, &compute.remote_ext_config).await {
+            match extension_server::download_extension(
+                &compute.remote_storage,
+                ExtensionType::Shared,
+            )
+            .await
+            {
                 Ok(_) => Response::new(Body::from("OK")),
                 Err(e) => {
                     error!("download_file failed: {}", e);