From 02f8650111ee07eb9cf93d20ed737930071c2111 Mon Sep 17 00:00:00 2001
From: Anastasia Lubennikova <anastasia@neon.tech>
Date: Thu, 8 Feb 2024 01:40:05 +0000
Subject: [PATCH] Add feature flag RemoteExtensionsUseLatest to compute_ctl.
 This will allow us to test new versions of extensions, without waiting for
 main branch commit

---
 compute_tools/src/compute.rs  | 11 ++++++++++-
 compute_tools/src/http/api.rs | 17 +++++++++++------
 libs/compute_api/src/spec.rs  |  9 +++++++++
 3 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index 0ca1a47fbf..76e600d768 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -1235,10 +1235,19 @@ LIMIT 100",
 
         info!("Downloading to shared preload libraries: {:?}", &libs_vec);
 
+        let build_tag_str = if spec
+            .features
+            .contains(&ComputeFeature::RemoteExtensionsUseLatest)
+        {
+            "latest"
+        } else {
+            &self.build_tag
+        };
+
         let mut download_tasks = Vec::new();
         for library in &libs_vec {
             let (ext_name, ext_path) =
-                remote_extensions.get_ext(library, true, &self.build_tag, &self.pgversion)?;
+                remote_extensions.get_ext(library, true, build_tag_str, &self.pgversion)?;
             download_tasks.push(self.download_extension(ext_name, ext_path));
         }
         let results = join_all(download_tasks).await;
diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs
index fa2c4cff28..765dab2245 100644
--- a/compute_tools/src/http/api.rs
+++ b/compute_tools/src/http/api.rs
@@ -8,6 +8,7 @@ use std::thread;
 use crate::compute::{ComputeNode, ComputeState, ParsedSpec};
 use compute_api::requests::ConfigurationRequest;
 use compute_api::responses::{ComputeStatus, ComputeStatusResponse, GenericAPIError};
+use compute_api::spec::ComputeFeature;
 
 use anyhow::Result;
 use hyper::service::{make_service_fn, service_fn};
@@ -171,12 +172,16 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
                     }
                 };
 
-                remote_extensions.get_ext(
-                    &filename,
-                    is_library,
-                    &compute.build_tag,
-                    &compute.pgversion,
-                )
+                let build_tag_str = if spec
+                    .features
+                    .contains(&ComputeFeature::RemoteExtensionsUseLatest)
+                {
+                    "latest"
+                } else {
+                    &compute.build_tag
+                };
+
+                remote_extensions.get_ext(&filename, is_library, build_tag_str, &compute.pgversion)
             };
 
             match ext {
diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs
index 13ac18e0c5..1f13a88be4 100644
--- a/libs/compute_api/src/spec.rs
+++ b/libs/compute_api/src/spec.rs
@@ -90,6 +90,11 @@ pub enum ComputeFeature {
     /// track short-lived connections as user activity.
     ActivityMonitorExperimental,
 
+    // Use latest version of remote extensions
+    // This is needed to allow us to test new versions of extensions before
+    // they are merged into the main branch.
+    RemoteExtensionsUseLatest,
+
     /// This is a special feature flag that is used to represent unknown feature flags.
     /// Basically all unknown to enum flags are represented as this one. See unit test
     /// `parse_unknown_features()` for more details.
@@ -152,8 +157,12 @@ impl RemoteExtSpec {
                 //
                 // Keep it in sync with path generation in
                 // https://github.com/neondatabase/build-custom-extensions/tree/main
+                //
+                // if ComputeFeature::RemoteExtensionsUseLatest is enabled
+                // use "latest" as the build_tag
                 let archive_path_str =
                     format!("{build_tag}/{pg_major_version}/extensions/{real_ext_name}.tar.zst");
+
                 Ok((
                     real_ext_name.to_string(),
                     RemotePath::from_string(&archive_path_str)?,