From d2e0d23159487b5682c3ecae73207baab1744736 Mon Sep 17 00:00:00 2001
From: Sasha Krassovsky <sasha@neon.tech>
Date: Thu, 18 Jan 2024 02:30:06 +0000
Subject: [PATCH] Gate it behind feature flags

---
 compute_tools/src/compute.rs           | 10 ++++++----
 control_plane/src/endpoint.rs          | 11 +++++++++--
 libs/compute_api/src/spec.rs           |  3 +++
 test_runner/fixtures/neon_fixtures.py  |  1 +
 test_runner/regress/test_migrations.py |  2 +-
 5 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index 4ea6bf9987cd..07e0abe6ff0d 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -744,10 +744,12 @@ impl ComputeNode {
         // 'Close' connection
         drop(client);
 
-        thread::spawn(move || {
-            let mut client = Client::connect(connstr.as_str(), NoTls)?;
-            handle_migrations(&mut client)
-        });
+        if self.has_feature(ComputeFeature::Migrations) {
+            thread::spawn(move || {
+                let mut client = Client::connect(connstr.as_str(), NoTls)?;
+                handle_migrations(&mut client)
+            });
+        }
         Ok(())
     }
 
diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs
index 43f8ea3b43d5..d3b0366d313e 100644
--- a/control_plane/src/endpoint.rs
+++ b/control_plane/src/endpoint.rs
@@ -57,7 +57,7 @@ use crate::local_env::LocalEnv;
 use crate::postgresql_conf::PostgresConf;
 
 use compute_api::responses::{ComputeState, ComputeStatus};
-use compute_api::spec::{Cluster, ComputeMode, ComputeSpec};
+use compute_api::spec::{Cluster, ComputeFeature, ComputeMode, ComputeSpec};
 
 // contents of a endpoint.json file
 #[derive(Serialize, Deserialize, PartialEq, Eq, Clone, Debug)]
@@ -70,6 +70,7 @@ pub struct EndpointConf {
     http_port: u16,
     pg_version: u32,
     skip_pg_catalog_updates: bool,
+    features: Vec<ComputeFeature>,
 }
 
 //
@@ -140,6 +141,7 @@ impl ComputeControlPlane {
             // with this we basically test a case of waking up an idle compute, where
             // we also skip catalog updates in the cloud.
             skip_pg_catalog_updates: true,
+            features: vec![],
         });
 
         ep.create_endpoint_dir()?;
@@ -154,6 +156,7 @@ impl ComputeControlPlane {
                 pg_port,
                 pg_version,
                 skip_pg_catalog_updates: true,
+                features: vec![],
             })?,
         )?;
         std::fs::write(
@@ -215,6 +218,9 @@ pub struct Endpoint {
 
     // Optimizations
     skip_pg_catalog_updates: bool,
+
+    // Feature flags
+    features: Vec<ComputeFeature>,
 }
 
 impl Endpoint {
@@ -244,6 +250,7 @@ impl Endpoint {
             tenant_id: conf.tenant_id,
             pg_version: conf.pg_version,
             skip_pg_catalog_updates: conf.skip_pg_catalog_updates,
+            features: conf.features,
         })
     }
 
@@ -519,7 +526,7 @@ impl Endpoint {
             skip_pg_catalog_updates: self.skip_pg_catalog_updates,
             format_version: 1.0,
             operation_uuid: None,
-            features: vec![],
+            features: self.features.clone(),
             cluster: Cluster {
                 cluster_id: None, // project ID: not used
                 name: None,       // project name: not used
diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs
index 13ac18e0c51d..5361d1400449 100644
--- a/libs/compute_api/src/spec.rs
+++ b/libs/compute_api/src/spec.rs
@@ -90,6 +90,9 @@ pub enum ComputeFeature {
     /// track short-lived connections as user activity.
     ActivityMonitorExperimental,
 
+    /// Enable running migrations
+    Migrations,
+
     /// 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.
diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
index ffd93004d2b3..65224bb5081d 100644
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -2914,6 +2914,7 @@ def respec(self, **kwargs):
 
         # Write it back updated
         with open(config_path, "w") as file:
+            log.info(json.dumps(dict(data_dict, **kwargs)))
             json.dump(dict(data_dict, **kwargs), file, indent=4)
 
     # Mock the extension part of spec passed from control plane for local testing
diff --git a/test_runner/regress/test_migrations.py b/test_runner/regress/test_migrations.py
index 51079e243e95..121fa91f66bc 100644
--- a/test_runner/regress/test_migrations.py
+++ b/test_runner/regress/test_migrations.py
@@ -10,7 +10,7 @@ def test_migrations(neon_simple_env: NeonEnv):
     endpoint = env.endpoints.create("test_migrations")
     log_path = endpoint.endpoint_path() / "compute.log"
 
-    endpoint.respec(skip_pg_catalog_updates=False)
+    endpoint.respec(skip_pg_catalog_updates=False, features=["migrations"])
     endpoint.start()
 
     time.sleep(1)  # Sleep to let migrations run