From a40ed86d874a03a5346ab85921f55d45e58c7cfb 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           | 43 +++++++++++++++++++++++++++--
 pgxn/neon/control_plane_connector.c |  2 +-
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs
index 09b8b19d10..005e3334b0 100644
--- a/compute_tools/src/spec.rs
+++ b/compute_tools/src/spec.rs
@@ -731,9 +731,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 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;
 
-    let mut query = "CREATE SCHEMA IF NOT EXISTS neon_migration AUTHORIZATION cloud_admin";
+    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";
@@ -745,6 +779,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])?;
@@ -755,5 +790,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 e467a9c43a..f6f006cba4 100644
--- a/pgxn/neon/control_plane_connector.c
+++ b/pgxn/neon/control_plane_connector.c
@@ -637,7 +637,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)