From c4c6158f7baad7b9145bf565c59fd019e3883e6c Mon Sep 17 00:00:00 2001
From: Sasha Krassovsky <krassovskysasha@gmail.com>
Date: Mon, 11 Dec 2023 12:06:29 -0800
Subject: [PATCH] Add test for migrations, add initial migration

---
 compute_tools/src/spec.rs           | 45 +++++++++++++++++++++++++++--
 pgxn/neon/control_plane_connector.c |  2 +-
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs
index b56ff60c6f7f0..b9c87add5d36e 100644
--- a/compute_tools/src/spec.rs
+++ b/compute_tools/src/spec.rs
@@ -701,9 +701,43 @@ pub fn handle_extension_neon(client: &mut Client) -> Result<()> {
 #[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";
+    let migrations = vec![
+        "ALTER ROLE neon_superuser BYPASSRLS",
+        r#"
+DO $$
+DECLARE
+    role_name text;
+BEGIN
+    FOR role_name IN SELECT rolname FROM pg_roles WHERE pg_has_role(rolname, 'neon_superuser', 'member')
+    LOOP
+        EXECUTE 'ALTER ROLE ' || quote_ident(role_name) || ' INHERIT';
+    END LOOP;
+
+    FOR role_name IN SELECT rolname FROM pg_roles
+        WHERE
+            NOT pg_has_role(rolname, 'neon_superuser', 'member') AND
+            rolname NOT IN ('pg_read_all_data',
+                            'pg_write_all_data',
+                            'pg_read_all_settings',
+                            'pg_read_all_stats',
+                            'pg_stat_scan_tables',
+                            'pg_monitor',
+                            'pg_database_owner',
+                            'pg_signal_backend',
+                            'pg_read_server_files',
+                            'pg_write_server_files',
+                            'pg_execute_server_program',
+                            'pg_checkpoint',
+                            'pg_use_reserved_connections',
+                            'pg_create_subscription')
+    LOOP
+        EXECUTE 'ALTER ROLE ' || quote_ident(role_name) || ' NOBYPASSRLS';
+    END LOOP;
+END $$;
+"#,
+    ];
+
+    let mut query = "CREATE SCHEMA IF NOT EXISTS neon_migration";
     client.simple_query(query)?;
 
     query = "CREATE SEQUENCE IF NOT EXISTS neon_migration.migration_id START WITH 0 INCREMENT BY 1 MINVALUE 0";
@@ -715,6 +749,7 @@ pub fn handle_migrations(client: &mut Client) -> Result<()> {
     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;
+    let starting_migration_id = current_migration;
 
     while current_migration < migrations.len() {
         client.simple_query(migrations[current_migration])?;
@@ -725,5 +760,9 @@ pub fn handle_migrations(client: &mut Client) -> Result<()> {
         migrations.len()
     );
     client.simple_query(&setval)?;
+    info!(
+        "Ran {} migrations",
+        (migrations.len() - starting_migration_id)
+    );
     Ok(())
 }
diff --git a/pgxn/neon/control_plane_connector.c b/pgxn/neon/control_plane_connector.c
index 2e7da671f9e8c..0e0fc7428c506 100644
--- a/pgxn/neon/control_plane_connector.c
+++ b/pgxn/neon/control_plane_connector.c
@@ -636,7 +636,7 @@ HandleAlterRole(AlterRoleStmt *stmt)
 	ListCell   *option;
 	const char *role_name = stmt->role->rolename;
 
-	if (RoleIsNeonSuperuser(role_name))
+	if (RoleIsNeonSuperuser(role_name) && !superuser())
 		elog(ERROR, "can't ALTER neon_superuser");
 
 	foreach(option, stmt->options)