From 89b8ea132ea0028b0ae5648b2e0f263bce7cbd68 Mon Sep 17 00:00:00 2001
From: Alek Westover <alek.westover@gmail.com>
Date: Wed, 21 Jun 2023 11:39:24 -0400
Subject: [PATCH] refactor more

---
 ALEK_LIST_FILES.txt                   |  2 +-
 compute_tools/src/extension_server.rs | 15 +++++++--------
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/ALEK_LIST_FILES.txt b/ALEK_LIST_FILES.txt
index a1419d1429..17b430abb5 100644
--- a/ALEK_LIST_FILES.txt
+++ b/ALEK_LIST_FILES.txt
@@ -1 +1 @@
-[RemotePath("tenants/9a8ce821f7946ed2f2d58f51f2595024/timelines/62cdc8653864e444171faac9a5c3cea9/000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__000000000169B098-000000000169B111"), RemotePath("tenants/9a8ce821f7946ed2f2d58f51f2595024/timelines/62cdc8653864e444171faac9a5c3cea9/index_part.json"), RemotePath("tenants/a9982e09ea4a00aff9c61daf12744098/timelines/55efa350b38ff1a9df45726dbaadbb9f/000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__000000000169B098-000000000169B111"), RemotePath("tenants/a9982e09ea4a00aff9c61daf12744098/timelines/55efa350b38ff1a9df45726dbaadbb9f/index_part.json"), RemotePath("tenants/a9982e09ea4a00aff9c61daf12744098/timelines/a0af9d76650e3477b7bd1a1e8e2793bf/000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__000000000169B098-000000000169B111"), RemotePath("tenants/a9982e09ea4a00aff9c61daf12744098/timelines/a0af9d76650e3477b7bd1a1e8e2793bf/index_part.json"), RemotePath("v15/share/extension/test_ext.control")]
\ No newline at end of file
+[RemotePath("tenants/75b0fcfc4874117b656b89640d29bcc1/timelines/b67a8af543c4d5e222f98345b8028a14/000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__000000000169B098-000000000169B111"), RemotePath("tenants/75b0fcfc4874117b656b89640d29bcc1/timelines/b67a8af543c4d5e222f98345b8028a14/index_part.json"), RemotePath("tenants/b2fbcd019b422e552505cbeab214eebf/timelines/27228ae7a75f4cefbf08ab17cf591030/000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__000000000169B098-000000000169B111"), RemotePath("tenants/b2fbcd019b422e552505cbeab214eebf/timelines/27228ae7a75f4cefbf08ab17cf591030/index_part.json"), RemotePath("tenants/b2fbcd019b422e552505cbeab214eebf/timelines/65e175c0d402139051ebf7353ddc737c/000000000000000000000000000000000000-FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF__000000000169B098-000000000169B111"), RemotePath("tenants/b2fbcd019b422e552505cbeab214eebf/timelines/65e175c0d402139051ebf7353ddc737c/index_part.json"), RemotePath("v15/share/extension/test_ext.control")]
\ No newline at end of file
diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs
index 4874e8adab..f1ff9034f7 100644
--- a/compute_tools/src/extension_server.rs
+++ b/compute_tools/src/extension_server.rs
@@ -11,8 +11,8 @@ 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 s3_config = create_s3_config(remote_ext_config)?;
-    download_extension(&s3_config, ExtensionType::Shared).await?;
+    let remote_storage = init_remote_storage(remote_ext_config)?;
+    download_extension(&remote_storage, ExtensionType::Shared).await?;
     Ok(())
 }
 
@@ -55,11 +55,9 @@ pub enum ExtensionType {
 }
 
 pub async fn download_extension(
-    config: &RemoteStorageConfig,
+    remote_storage: &GenericRemoteStorage,
     ext_type: ExtensionType,
 ) -> anyhow::Result<()> {
-    let remote_storage = GenericRemoteStorage::from_config(config)?;
-
     let from_paths = remote_storage.list_files(None).await?;
     std::fs::write("ALEK_LIST_FILES.txt", format!("{:?}", from_paths))?;
 
@@ -103,7 +101,7 @@ pub async fn download_extension(
     Ok(())
 }
 
-pub fn create_s3_config(remote_ext_config: &str) -> anyhow::Result<RemoteStorageConfig> {
+pub fn init_remote_storage(remote_ext_config: &str) -> anyhow::Result<GenericRemoteStorage> {
     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,
@@ -127,9 +125,10 @@ pub fn create_s3_config(remote_ext_config: &str) -> anyhow::Result<RemoteStorage
         concurrency_limit: NonZeroUsize::new(100).expect("100 != 0"),
         max_keys_per_list_response: None,
     };
-    Ok(RemoteStorageConfig {
+    let config = RemoteStorageConfig {
         max_concurrent_syncs: NonZeroUsize::new(100).expect("100 != 0"),
         max_sync_errors: NonZeroU32::new(100).expect("100 != 0"),
         storage: RemoteStorageKind::AwsS3(config),
-    })
+    };
+    Ok(GenericRemoteStorage::from_config(&config)?)
 }