From 9d3c573fbef060a15f8093588146128a4f236f5d Mon Sep 17 00:00:00 2001
From: Sasha Krassovsky <krassovskysasha@gmail.com>
Date: Fri, 8 Dec 2023 13:18:35 -0800
Subject: [PATCH] Add support for migrations within compute_ctl

---
 compute_tools/src/compute.rs |  5 +++++
 compute_tools/src/spec.rs    | 30 ++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index b39a800f1450..8da91615c96e 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -710,6 +710,7 @@ impl ComputeNode {
         handle_grants(spec, &mut client, self.connstr.as_str())?;
         handle_extensions(spec, &mut client)?;
         handle_extension_neon(&mut client)?;
+        handle_migrations(&mut client)?;
         create_availability_check_data(&mut client)?;
 
         // 'Close' connection
@@ -760,6 +761,10 @@ impl ComputeNode {
             handle_grants(&spec, &mut client, self.connstr.as_str())?;
             handle_extensions(&spec, &mut client)?;
             handle_extension_neon(&mut client)?;
+            // We can skip handle_migrations here because a new migration can only appear
+            // if we have a new version of the compute_ctl binary, which can only happen
+            // if compute got restarted, in which case we'll end up inside of apply_config
+            // instead of reconfigure.
         }
 
         // 'Close' connection
diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs
index ba1ee6d1b217..b56ff60c6f7f 100644
--- a/compute_tools/src/spec.rs
+++ b/compute_tools/src/spec.rs
@@ -697,3 +697,33 @@ pub fn handle_extension_neon(client: &mut Client) -> Result<()> {
 
     Ok(())
 }
+
+#[instrument(skip_all)]
+pub fn handle_migrations(client: &mut Client) -> Result<()> {
+    info!("handle migrations");
+    let migrations = vec![""];
+
+    let mut query = "CREATE SCHEMA IF NOT EXISTS neon_migration AUTHORIZATION cloud_admin";
+    client.simple_query(query)?;
+
+    query = "CREATE SEQUENCE IF NOT EXISTS neon_migration.migration_id START WITH 0 INCREMENT BY 1 MINVALUE 0";
+    client.simple_query(query)?;
+
+    query = "GRANT USAGE, SELECT ON SEQUENCE neon_migration.migration_id TO cloud_admin";
+    client.simple_query(query)?;
+
+    query = "SELECT last_value FROM neon_migration.migration_id";
+    let row = client.query_one(query, &[])?;
+    let mut current_migration: usize = row.get::<&str, i64>("last_value") as usize;
+
+    while current_migration < migrations.len() {
+        client.simple_query(migrations[current_migration])?;
+        current_migration += 1;
+    }
+    let setval = format!(
+        "SELECT setval('neon_migration.migration_id', {})",
+        migrations.len()
+    );
+    client.simple_query(&setval)?;
+    Ok(())
+}