LCOV - code coverage report
Current view: top level - compute_tools/src - spec_apply.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 0.0 % 874 0
Test Date: 2025-03-12 00:01:28 Functions: 0.0 % 55 0

            Line data    Source code
       1              : use std::collections::{HashMap, HashSet};
       2              : use std::fmt::{Debug, Formatter};
       3              : use std::future::Future;
       4              : use std::iter::{empty, once};
       5              : use std::sync::Arc;
       6              : 
       7              : use anyhow::{Context, Result};
       8              : use compute_api::responses::ComputeStatus;
       9              : use compute_api::spec::{ComputeAudit, ComputeFeature, ComputeSpec, Database, PgIdent, Role};
      10              : use futures::future::join_all;
      11              : use tokio::sync::RwLock;
      12              : use tokio_postgres::Client;
      13              : use tokio_postgres::error::SqlState;
      14              : use tracing::{Instrument, debug, error, info, info_span, instrument, warn};
      15              : 
      16              : use crate::compute::{ComputeNode, ComputeState};
      17              : use crate::pg_helpers::{
      18              :     DatabaseExt, Escaping, GenericOptionsSearch, RoleExt, get_existing_dbs_async,
      19              :     get_existing_roles_async,
      20              : };
      21              : use crate::spec_apply::ApplySpecPhase::{
      22              :     CreateAndAlterDatabases, CreateAndAlterRoles, CreateAvailabilityCheck, CreateNeonSuperuser,
      23              :     CreatePgauditExtension, CreatePgauditlogtofileExtension, CreateSchemaNeon,
      24              :     DisablePostgresDBPgAudit, DropInvalidDatabases, DropRoles, FinalizeDropLogicalSubscriptions,
      25              :     HandleNeonExtension, HandleOtherExtensions, RenameAndDeleteDatabases, RenameRoles,
      26              :     RunInEachDatabase,
      27              : };
      28              : use crate::spec_apply::PerDatabasePhase::{
      29              :     ChangeSchemaPerms, DeleteDBRoleReferences, DropLogicalSubscriptions, HandleAnonExtension,
      30              : };
      31              : 
      32              : impl ComputeNode {
      33              :     /// Apply the spec to the running PostgreSQL instance.
      34              :     /// The caller can decide to run with multiple clients in parallel, or
      35              :     /// single mode.  Either way, the commands executed will be the same, and
      36              :     /// only commands run in different databases are parallelized.
      37              :     #[instrument(skip_all)]
      38              :     pub fn apply_spec_sql(
      39              :         &self,
      40              :         spec: Arc<ComputeSpec>,
      41              :         conf: Arc<tokio_postgres::Config>,
      42              :         concurrency: usize,
      43              :     ) -> Result<()> {
      44              :         info!("Applying config with max {} concurrency", concurrency);
      45              :         debug!("Config: {:?}", spec);
      46              : 
      47              :         let rt = tokio::runtime::Handle::current();
      48            0 :         rt.block_on(async {
      49              :             // Proceed with post-startup configuration. Note, that order of operations is important.
      50            0 :             let client = Self::get_maintenance_client(&conf).await?;
      51            0 :             let spec = spec.clone();
      52              : 
      53            0 :             let databases = get_existing_dbs_async(&client).await?;
      54            0 :             let roles = get_existing_roles_async(&client)
      55            0 :                 .await?
      56            0 :                 .into_iter()
      57            0 :                 .map(|role| (role.name.clone(), role))
      58            0 :                 .collect::<HashMap<String, Role>>();
      59            0 : 
      60            0 :             // Check if we need to drop subscriptions before starting the endpoint.
      61            0 :             //
      62            0 :             // It is important to do this operation exactly once when endpoint starts on a new branch.
      63            0 :             // Otherwise, we may drop not inherited, but newly created subscriptions.
      64            0 :             //
      65            0 :             // We cannot rely only on spec.drop_subscriptions_before_start flag,
      66            0 :             // because if for some reason compute restarts inside VM,
      67            0 :             // it will start again with the same spec and flag value.
      68            0 :             //
      69            0 :             // To handle this, we save the fact of the operation in the database
      70            0 :             // in the neon.drop_subscriptions_done table.
      71            0 :             // If the table does not exist, we assume that the operation was never performed, so we must do it.
      72            0 :             // If table exists, we check if the operation was performed on the current timelilne.
      73            0 :             //
      74            0 :             let mut drop_subscriptions_done = false;
      75            0 : 
      76            0 :             if spec.drop_subscriptions_before_start {
      77            0 :                 let timeline_id = self.get_timeline_id().context("timeline_id must be set")?;
      78            0 :                 let query = format!("select 1 from neon.drop_subscriptions_done where timeline_id = '{}'", timeline_id);
      79            0 : 
      80            0 :                 info!("Checking if drop subscription operation was already performed for timeline_id: {}", timeline_id);
      81              : 
      82              :                 drop_subscriptions_done =  match
      83            0 :                     client.simple_query(&query).await {
      84            0 :                     Ok(result) => {
      85            0 :                         matches!(&result[0], postgres::SimpleQueryMessage::Row(_))
      86              :                     },
      87            0 :                     Err(e) =>
      88            0 :                     {
      89            0 :                         match e.code() {
      90            0 :                             Some(&SqlState::UNDEFINED_TABLE) => false,
      91              :                             _ => {
      92              :                                 // We don't expect any other error here, except for the schema/table not existing
      93            0 :                                 error!("Error checking if drop subscription operation was already performed: {}", e);
      94            0 :                                 return Err(e.into());
      95              :                             }
      96              :                         }
      97              :                     }
      98              :                 }
      99            0 :             };
     100              : 
     101              : 
     102            0 :             let jwks_roles = Arc::new(
     103            0 :                 spec.as_ref()
     104            0 :                     .local_proxy_config
     105            0 :                     .iter()
     106            0 :                     .flat_map(|it| &it.jwks)
     107            0 :                     .flatten()
     108            0 :                     .flat_map(|setting| &setting.role_names)
     109            0 :                     .cloned()
     110            0 :                     .collect::<HashSet<_>>(),
     111            0 :             );
     112            0 : 
     113            0 :             let ctx = Arc::new(tokio::sync::RwLock::new(MutableApplyContext {
     114            0 :                 roles,
     115            0 :                 dbs: databases,
     116            0 :             }));
     117            0 : 
     118            0 :             // Apply special pre drop database phase.
     119            0 :             // NOTE: we use the code of RunInEachDatabase phase for parallelism
     120            0 :             // and connection management, but we don't really run it in *each* database,
     121            0 :             // only in databases, we're about to drop.
     122            0 :             info!("Applying PerDatabase (pre-dropdb) phase");
     123            0 :             let concurrency_token = Arc::new(tokio::sync::Semaphore::new(concurrency));
     124            0 : 
     125            0 :             // Run the phase for each database that we're about to drop.
     126            0 :             let db_processes = spec
     127            0 :                 .delta_operations
     128            0 :                 .iter()
     129            0 :                 .flatten()
     130            0 :                 .filter_map(move |op| {
     131            0 :                     if op.action.as_str() == "delete_db" {
     132            0 :                         Some(op.name.clone())
     133              :                     } else {
     134            0 :                         None
     135              :                     }
     136            0 :                 })
     137            0 :                 .map(|dbname| {
     138            0 :                     let spec = spec.clone();
     139            0 :                     let ctx = ctx.clone();
     140            0 :                     let jwks_roles = jwks_roles.clone();
     141            0 :                     let mut conf = conf.as_ref().clone();
     142            0 :                     let concurrency_token = concurrency_token.clone();
     143            0 :                     // We only need dbname field for this phase, so set other fields to dummy values
     144            0 :                     let db = DB::UserDB(Database {
     145            0 :                         name: dbname.clone(),
     146            0 :                         owner: "cloud_admin".to_string(),
     147            0 :                         options: None,
     148            0 :                         restrict_conn: false,
     149            0 :                         invalid: false,
     150            0 :                     });
     151            0 : 
     152            0 :                     debug!("Applying per-database phases for Database {:?}", &db);
     153              : 
     154            0 :                     match &db {
     155            0 :                         DB::SystemDB => {}
     156            0 :                         DB::UserDB(db) => {
     157            0 :                             conf.dbname(db.name.as_str());
     158            0 :                         }
     159              :                     }
     160              : 
     161            0 :                     let conf = Arc::new(conf);
     162            0 :                     let fut = Self::apply_spec_sql_db(
     163            0 :                         spec.clone(),
     164            0 :                         conf,
     165            0 :                         ctx.clone(),
     166            0 :                         jwks_roles.clone(),
     167            0 :                         concurrency_token.clone(),
     168            0 :                         db,
     169            0 :                         [DropLogicalSubscriptions].to_vec(),
     170            0 :                     );
     171            0 : 
     172            0 :                     Ok(tokio::spawn(fut))
     173            0 :                 })
     174            0 :                 .collect::<Vec<Result<_, anyhow::Error>>>();
     175              : 
     176            0 :             for process in db_processes.into_iter() {
     177            0 :                 let handle = process?;
     178            0 :                 if let Err(e) = handle.await? {
     179              :                     // Handle the error case where the database does not exist
     180              :                     // We do not check whether the DB exists or not in the deletion phase,
     181              :                     // so we shouldn't be strict about it in pre-deletion cleanup as well.
     182            0 :                     if e.to_string().contains("does not exist") {
     183            0 :                         warn!("Error dropping subscription: {}", e);
     184              :                     } else {
     185            0 :                         return Err(e);
     186              :                     }
     187            0 :                 };
     188              :             }
     189              : 
     190            0 :             for phase in [
     191            0 :                 CreateNeonSuperuser,
     192            0 :                 DropInvalidDatabases,
     193            0 :                 RenameRoles,
     194            0 :                 CreateAndAlterRoles,
     195            0 :                 RenameAndDeleteDatabases,
     196            0 :                 CreateAndAlterDatabases,
     197            0 :                 CreateSchemaNeon,
     198              :             ] {
     199            0 :                 info!("Applying phase {:?}", &phase);
     200            0 :                 apply_operations(
     201            0 :                     spec.clone(),
     202            0 :                     ctx.clone(),
     203            0 :                     jwks_roles.clone(),
     204            0 :                     phase,
     205            0 :                     || async { Ok(&client) },
     206            0 :                 )
     207            0 :                 .await?;
     208              :             }
     209              : 
     210            0 :             info!("Applying RunInEachDatabase2 phase");
     211            0 :             let concurrency_token = Arc::new(tokio::sync::Semaphore::new(concurrency));
     212            0 : 
     213            0 :             let db_processes = spec
     214            0 :                 .cluster
     215            0 :                 .databases
     216            0 :                 .iter()
     217            0 :                 .map(|db| DB::new(db.clone()))
     218            0 :                 // include
     219            0 :                 .chain(once(DB::SystemDB))
     220            0 :                 .map(|db| {
     221            0 :                     let spec = spec.clone();
     222            0 :                     let ctx = ctx.clone();
     223            0 :                     let jwks_roles = jwks_roles.clone();
     224            0 :                     let mut conf = conf.as_ref().clone();
     225            0 :                     let concurrency_token = concurrency_token.clone();
     226            0 :                     let db = db.clone();
     227            0 : 
     228            0 :                     debug!("Applying per-database phases for Database {:?}", &db);
     229              : 
     230            0 :                     match &db {
     231            0 :                         DB::SystemDB => {}
     232            0 :                         DB::UserDB(db) => {
     233            0 :                             conf.dbname(db.name.as_str());
     234            0 :                         }
     235              :                     }
     236              : 
     237            0 :                     let conf = Arc::new(conf);
     238            0 :                     let mut phases = vec![
     239            0 :                         DeleteDBRoleReferences,
     240            0 :                         ChangeSchemaPerms,
     241            0 :                         HandleAnonExtension,
     242            0 :                     ];
     243            0 : 
     244            0 :                     if spec.drop_subscriptions_before_start && !drop_subscriptions_done {
     245            0 :                         info!("Adding DropLogicalSubscriptions phase because drop_subscriptions_before_start is set");
     246            0 :                         phases.push(DropLogicalSubscriptions);
     247            0 :                     }
     248              : 
     249            0 :                     let fut = Self::apply_spec_sql_db(
     250            0 :                         spec.clone(),
     251            0 :                         conf,
     252            0 :                         ctx.clone(),
     253            0 :                         jwks_roles.clone(),
     254            0 :                         concurrency_token.clone(),
     255            0 :                         db,
     256            0 :                         phases,
     257            0 :                     );
     258            0 : 
     259            0 :                     Ok(tokio::spawn(fut))
     260            0 :                 })
     261            0 :                 .collect::<Vec<Result<_, anyhow::Error>>>();
     262              : 
     263            0 :             for process in db_processes.into_iter() {
     264            0 :                 let handle = process?;
     265            0 :                 handle.await??;
     266              :             }
     267              : 
     268            0 :             let mut phases = vec![
     269            0 :                 HandleOtherExtensions,
     270            0 :                 HandleNeonExtension, // This step depends on CreateSchemaNeon
     271            0 :                 CreateAvailabilityCheck,
     272            0 :                 DropRoles,
     273            0 :             ];
     274            0 : 
     275            0 :             // This step depends on CreateSchemaNeon
     276            0 :             if spec.drop_subscriptions_before_start && !drop_subscriptions_done {
     277            0 :                 info!("Adding FinalizeDropLogicalSubscriptions phase because drop_subscriptions_before_start is set");
     278            0 :                 phases.push(FinalizeDropLogicalSubscriptions);
     279            0 :             }
     280              : 
     281              :             // Keep DisablePostgresDBPgAudit phase at the end,
     282              :             // so that all config operations are audit logged.
     283            0 :             match spec.audit_log_level
     284              :             {
     285            0 :                 ComputeAudit::Hipaa => {
     286            0 :                     phases.push(CreatePgauditExtension);
     287            0 :                     phases.push(CreatePgauditlogtofileExtension);
     288            0 :                     phases.push(DisablePostgresDBPgAudit);
     289            0 :                 }
     290            0 :                 ComputeAudit::Log => { /* not implemented yet */ }
     291            0 :                 ComputeAudit::Disabled => {}
     292              :             }
     293              : 
     294            0 :             for phase in phases {
     295            0 :                 debug!("Applying phase {:?}", &phase);
     296            0 :                 apply_operations(
     297            0 :                     spec.clone(),
     298            0 :                     ctx.clone(),
     299            0 :                     jwks_roles.clone(),
     300            0 :                     phase,
     301            0 :                     || async { Ok(&client) },
     302            0 :                 )
     303            0 :                 .await?;
     304              :             }
     305              : 
     306            0 :             Ok::<(), anyhow::Error>(())
     307            0 :         })?;
     308              : 
     309              :         Ok(())
     310              :     }
     311              : 
     312              :     /// Apply SQL migrations of the RunInEachDatabase phase.
     313              :     ///
     314              :     /// May opt to not connect to databases that don't have any scheduled
     315              :     /// operations.  The function is concurrency-controlled with the provided
     316              :     /// semaphore.  The caller has to make sure the semaphore isn't exhausted.
     317            0 :     async fn apply_spec_sql_db(
     318            0 :         spec: Arc<ComputeSpec>,
     319            0 :         conf: Arc<tokio_postgres::Config>,
     320            0 :         ctx: Arc<tokio::sync::RwLock<MutableApplyContext>>,
     321            0 :         jwks_roles: Arc<HashSet<String>>,
     322            0 :         concurrency_token: Arc<tokio::sync::Semaphore>,
     323            0 :         db: DB,
     324            0 :         subphases: Vec<PerDatabasePhase>,
     325            0 :     ) -> Result<()> {
     326            0 :         let _permit = concurrency_token.acquire().await?;
     327              : 
     328            0 :         let mut client_conn = None;
     329              : 
     330            0 :         for subphase in subphases {
     331            0 :             apply_operations(
     332            0 :                 spec.clone(),
     333            0 :                 ctx.clone(),
     334            0 :                 jwks_roles.clone(),
     335            0 :                 RunInEachDatabase {
     336            0 :                     db: db.clone(),
     337            0 :                     subphase,
     338            0 :                 },
     339            0 :                 // Only connect if apply_operation actually wants a connection.
     340            0 :                 // It's quite possible this database doesn't need any queries,
     341            0 :                 // so by not connecting we save time and effort connecting to
     342            0 :                 // that database.
     343            0 :                 || async {
     344            0 :                     if client_conn.is_none() {
     345            0 :                         let db_client = Self::get_maintenance_client(&conf).await?;
     346            0 :                         client_conn.replace(db_client);
     347            0 :                     }
     348            0 :                     let client = client_conn.as_ref().unwrap();
     349            0 :                     Ok(client)
     350            0 :                 },
     351            0 :             )
     352            0 :             .await?;
     353              :         }
     354              : 
     355            0 :         drop(client_conn);
     356            0 : 
     357            0 :         Ok::<(), anyhow::Error>(())
     358            0 :     }
     359              : 
     360              :     /// Choose how many concurrent connections to use for applying the spec changes.
     361            0 :     pub fn max_service_connections(
     362            0 :         &self,
     363            0 :         compute_state: &ComputeState,
     364            0 :         spec: &ComputeSpec,
     365            0 :     ) -> usize {
     366            0 :         // If the cluster is in Init state we don't have to deal with user connections,
     367            0 :         // and can thus use all `max_connections` connection slots. However, that's generally not
     368            0 :         // very efficient, so we generally still limit it to a smaller number.
     369            0 :         if compute_state.status == ComputeStatus::Init {
     370              :             // If the settings contain 'max_connections', use that as template
     371            0 :             if let Some(config) = spec.cluster.settings.find("max_connections") {
     372            0 :                 config.parse::<usize>().ok()
     373              :             } else {
     374              :                 // Otherwise, try to find the setting in the postgresql_conf string
     375            0 :                 spec.cluster
     376            0 :                     .postgresql_conf
     377            0 :                     .iter()
     378            0 :                     .flat_map(|conf| conf.split("\n"))
     379            0 :                     .filter_map(|line| {
     380            0 :                         if !line.contains("max_connections") {
     381            0 :                             return None;
     382            0 :                         }
     383              : 
     384            0 :                         let (key, value) = line.split_once("=")?;
     385            0 :                         let key = key
     386            0 :                             .trim_start_matches(char::is_whitespace)
     387            0 :                             .trim_end_matches(char::is_whitespace);
     388            0 : 
     389            0 :                         let value = value
     390            0 :                             .trim_start_matches(char::is_whitespace)
     391            0 :                             .trim_end_matches(char::is_whitespace);
     392            0 : 
     393            0 :                         if key != "max_connections" {
     394            0 :                             return None;
     395            0 :                         }
     396            0 : 
     397            0 :                         value.parse::<usize>().ok()
     398            0 :                     })
     399            0 :                     .next()
     400              :             }
     401              :             // If max_connections is present, use at most 1/3rd of that.
     402              :             // When max_connections is lower than 30, try to use at least 10 connections, but
     403              :             // never more than max_connections.
     404            0 :             .map(|limit| match limit {
     405            0 :                 0..10 => limit,
     406            0 :                 10..30 => 10,
     407            0 :                 30.. => limit / 3,
     408            0 :             })
     409            0 :             // If we didn't find max_connections, default to 10 concurrent connections.
     410            0 :             .unwrap_or(10)
     411              :         } else {
     412              :             // state == Running
     413              :             // Because the cluster is already in the Running state, we should assume users are
     414              :             // already connected to the cluster, and high concurrency could negatively
     415              :             // impact user connectivity. Therefore, we can limit concurrency to the number of
     416              :             // reserved superuser connections, which users wouldn't be able to use anyway.
     417            0 :             spec.cluster
     418            0 :                 .settings
     419            0 :                 .find("superuser_reserved_connections")
     420            0 :                 .iter()
     421            0 :                 .filter_map(|val| val.parse::<usize>().ok())
     422            0 :                 .map(|val| if val > 1 { val - 1 } else { 1 })
     423            0 :                 .last()
     424            0 :                 .unwrap_or(3)
     425              :         }
     426            0 :     }
     427              : }
     428              : 
     429              : #[derive(Clone)]
     430              : pub enum DB {
     431              :     SystemDB,
     432              :     UserDB(Database),
     433              : }
     434              : 
     435              : impl DB {
     436            0 :     pub fn new(db: Database) -> DB {
     437            0 :         Self::UserDB(db)
     438            0 :     }
     439              : 
     440            0 :     pub fn is_owned_by(&self, role: &PgIdent) -> bool {
     441            0 :         match self {
     442            0 :             DB::SystemDB => false,
     443            0 :             DB::UserDB(db) => &db.owner == role,
     444              :         }
     445            0 :     }
     446              : }
     447              : 
     448              : impl Debug for DB {
     449            0 :     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
     450            0 :         match self {
     451            0 :             DB::SystemDB => f.debug_tuple("SystemDB").finish(),
     452            0 :             DB::UserDB(db) => f.debug_tuple("UserDB").field(&db.name).finish(),
     453              :         }
     454            0 :     }
     455              : }
     456              : 
     457              : #[derive(Copy, Clone, Debug)]
     458              : pub enum PerDatabasePhase {
     459              :     DeleteDBRoleReferences,
     460              :     ChangeSchemaPerms,
     461              :     HandleAnonExtension,
     462              :     /// This is a shared phase, used for both i) dropping dangling LR subscriptions
     463              :     /// before dropping the DB, and ii) dropping all subscriptions after creating
     464              :     /// a fresh branch.
     465              :     /// N.B. we will skip all DBs that are not present in Postgres, invalid, or
     466              :     /// have `datallowconn = false` (`restrict_conn`).
     467              :     DropLogicalSubscriptions,
     468              : }
     469              : 
     470              : #[derive(Clone, Debug)]
     471              : pub enum ApplySpecPhase {
     472              :     CreateNeonSuperuser,
     473              :     DropInvalidDatabases,
     474              :     RenameRoles,
     475              :     CreateAndAlterRoles,
     476              :     RenameAndDeleteDatabases,
     477              :     CreateAndAlterDatabases,
     478              :     CreateSchemaNeon,
     479              :     RunInEachDatabase { db: DB, subphase: PerDatabasePhase },
     480              :     CreatePgauditExtension,
     481              :     CreatePgauditlogtofileExtension,
     482              :     DisablePostgresDBPgAudit,
     483              :     HandleOtherExtensions,
     484              :     HandleNeonExtension,
     485              :     CreateAvailabilityCheck,
     486              :     DropRoles,
     487              :     FinalizeDropLogicalSubscriptions,
     488              : }
     489              : 
     490              : pub struct Operation {
     491              :     pub query: String,
     492              :     pub comment: Option<String>,
     493              : }
     494              : 
     495              : pub struct MutableApplyContext {
     496              :     pub roles: HashMap<String, Role>,
     497              :     pub dbs: HashMap<String, Database>,
     498              : }
     499              : 
     500              : /// Apply the operations that belong to the given spec apply phase.
     501              : ///
     502              : /// Commands within a single phase are executed in order of Iterator yield.
     503              : /// Commands of ApplySpecPhase::RunInEachDatabase will execute in the database
     504              : /// indicated by its `db` field, and can share a single client for all changes
     505              : /// to that database.
     506              : ///
     507              : /// Notes:
     508              : /// - Commands are pipelined, and thus may cause incomplete apply if one
     509              : ///   command of many fails.
     510              : /// - Failing commands will fail the phase's apply step once the return value
     511              : ///   is processed.
     512              : /// - No timeouts have (yet) been implemented.
     513              : /// - The caller is responsible for limiting and/or applying concurrency.
     514            0 : pub async fn apply_operations<'a, Fut, F>(
     515            0 :     spec: Arc<ComputeSpec>,
     516            0 :     ctx: Arc<RwLock<MutableApplyContext>>,
     517            0 :     jwks_roles: Arc<HashSet<String>>,
     518            0 :     apply_spec_phase: ApplySpecPhase,
     519            0 :     client: F,
     520            0 : ) -> Result<()>
     521            0 : where
     522            0 :     F: FnOnce() -> Fut,
     523            0 :     Fut: Future<Output = Result<&'a Client>>,
     524            0 : {
     525            0 :     debug!("Starting phase {:?}", &apply_spec_phase);
     526            0 :     let span = info_span!("db_apply_changes", phase=?apply_spec_phase);
     527            0 :     let span2 = span.clone();
     528            0 :     async move {
     529            0 :         debug!("Processing phase {:?}", &apply_spec_phase);
     530            0 :         let ctx = ctx;
     531              : 
     532            0 :         let mut ops = get_operations(&spec, &ctx, &jwks_roles, &apply_spec_phase)
     533            0 :             .await?
     534            0 :             .peekable();
     535            0 : 
     536            0 :         // Return (and by doing so, skip requesting the PostgreSQL client) if
     537            0 :         // we don't have any operations scheduled.
     538            0 :         if ops.peek().is_none() {
     539            0 :             return Ok(());
     540            0 :         }
     541              : 
     542            0 :         let client = client().await?;
     543              : 
     544            0 :         debug!("Applying phase {:?}", &apply_spec_phase);
     545              : 
     546            0 :         let active_queries = ops
     547            0 :             .map(|op| {
     548            0 :                 let Operation { comment, query } = op;
     549            0 :                 let inspan = match comment {
     550            0 :                     None => span.clone(),
     551            0 :                     Some(comment) => info_span!("phase {}: {}", comment),
     552              :                 };
     553              : 
     554            0 :                 async {
     555            0 :                     let query = query;
     556            0 :                     let res = client.simple_query(&query).await;
     557            0 :                     debug!(
     558            0 :                         "{} {}",
     559            0 :                         if res.is_ok() {
     560            0 :                             "successfully executed"
     561              :                         } else {
     562            0 :                             "failed to execute"
     563              :                         },
     564              :                         query
     565              :                     );
     566            0 :                     res
     567            0 :                 }
     568            0 :                 .instrument(inspan)
     569            0 :             })
     570            0 :             .collect::<Vec<_>>();
     571            0 : 
     572            0 :         drop(ctx);
     573              : 
     574            0 :         for it in join_all(active_queries).await {
     575            0 :             drop(it?);
     576              :         }
     577              : 
     578            0 :         debug!("Completed phase {:?}", &apply_spec_phase);
     579              : 
     580            0 :         Ok(())
     581            0 :     }
     582            0 :     .instrument(span2)
     583            0 :     .await
     584            0 : }
     585              : 
     586              : /// Create a stream of operations to be executed for that phase of applying
     587              : /// changes.
     588              : ///
     589              : /// In the future we may generate a single stream of changes and then
     590              : /// sort/merge/batch execution, but for now this is a nice way to improve
     591              : /// batching behavior of the commands.
     592            0 : async fn get_operations<'a>(
     593            0 :     spec: &'a ComputeSpec,
     594            0 :     ctx: &'a RwLock<MutableApplyContext>,
     595            0 :     jwks_roles: &'a HashSet<String>,
     596            0 :     apply_spec_phase: &'a ApplySpecPhase,
     597            0 : ) -> Result<Box<dyn Iterator<Item = Operation> + 'a + Send>> {
     598            0 :     match apply_spec_phase {
     599            0 :         ApplySpecPhase::CreateNeonSuperuser => Ok(Box::new(once(Operation {
     600            0 :             query: include_str!("sql/create_neon_superuser.sql").to_string(),
     601            0 :             comment: None,
     602            0 :         }))),
     603              :         ApplySpecPhase::DropInvalidDatabases => {
     604            0 :             let mut ctx = ctx.write().await;
     605            0 :             let databases = &mut ctx.dbs;
     606            0 : 
     607            0 :             let keys: Vec<_> = databases
     608            0 :                 .iter()
     609            0 :                 .filter(|(_, db)| db.invalid)
     610            0 :                 .map(|(dbname, _)| dbname.clone())
     611            0 :                 .collect();
     612            0 : 
     613            0 :             // After recent commit in Postgres, interrupted DROP DATABASE
     614            0 :             // leaves the database in the invalid state. According to the
     615            0 :             // commit message, the only option for user is to drop it again.
     616            0 :             // See:
     617            0 :             //   https://github.com/postgres/postgres/commit/a4b4cc1d60f7e8ccfcc8ff8cb80c28ee411ad9a9
     618            0 :             //
     619            0 :             // Postgres Neon extension is done the way, that db is de-registered
     620            0 :             // in the control plane metadata only after it is dropped. So there is
     621            0 :             // a chance that it still thinks that the db should exist. This means
     622            0 :             // that it will be re-created by the `CreateDatabases` phase. This
     623            0 :             // is fine, as user can just drop the table again (in vanilla
     624            0 :             // Postgres they would need to do the same).
     625            0 :             let operations = keys
     626            0 :                 .into_iter()
     627            0 :                 .filter_map(move |dbname| ctx.dbs.remove(&dbname))
     628            0 :                 .map(|db| Operation {
     629            0 :                     query: format!("DROP DATABASE IF EXISTS {}", db.name.pg_quote()),
     630            0 :                     comment: Some(format!("Dropping invalid database {}", db.name)),
     631            0 :                 });
     632            0 : 
     633            0 :             Ok(Box::new(operations))
     634              :         }
     635              :         ApplySpecPhase::RenameRoles => {
     636            0 :             let mut ctx = ctx.write().await;
     637              : 
     638            0 :             let operations = spec
     639            0 :                 .delta_operations
     640            0 :                 .iter()
     641            0 :                 .flatten()
     642            0 :                 .filter(|op| op.action == "rename_role")
     643            0 :                 .filter_map(move |op| {
     644            0 :                     let roles = &mut ctx.roles;
     645            0 : 
     646            0 :                     if roles.contains_key(op.name.as_str()) {
     647            0 :                         None
     648              :                     } else {
     649            0 :                         let new_name = op.new_name.as_ref().unwrap();
     650            0 :                         let mut role = roles.remove(op.name.as_str()).unwrap();
     651            0 : 
     652            0 :                         role.name = new_name.clone();
     653            0 :                         role.encrypted_password = None;
     654            0 :                         roles.insert(role.name.clone(), role);
     655            0 : 
     656            0 :                         Some(Operation {
     657            0 :                             query: format!(
     658            0 :                                 "ALTER ROLE {} RENAME TO {}",
     659            0 :                                 op.name.pg_quote(),
     660            0 :                                 new_name.pg_quote()
     661            0 :                             ),
     662            0 :                             comment: Some(format!("renaming role '{}' to '{}'", op.name, new_name)),
     663            0 :                         })
     664              :                     }
     665            0 :                 });
     666            0 : 
     667            0 :             Ok(Box::new(operations))
     668              :         }
     669              :         ApplySpecPhase::CreateAndAlterRoles => {
     670            0 :             let mut ctx = ctx.write().await;
     671              : 
     672            0 :             let operations = spec.cluster.roles
     673            0 :                 .iter()
     674            0 :                 .filter_map(move |role| {
     675            0 :                     let roles = &mut ctx.roles;
     676            0 :                     let db_role = roles.get(&role.name);
     677            0 : 
     678            0 :                     match db_role {
     679            0 :                         Some(db_role) => {
     680            0 :                             if db_role.encrypted_password != role.encrypted_password {
     681              :                                 // This can be run on /every/ role! Not just ones created through the console.
     682              :                                 // This means that if you add some funny ALTER here that adds a permission,
     683              :                                 // this will get run even on user-created roles! This will result in different
     684              :                                 // behavior before and after a spec gets reapplied. The below ALTER as it stands
     685              :                                 // now only grants LOGIN and changes the password. Please do not allow this branch
     686              :                                 // to do anything silly.
     687            0 :                                 Some(Operation {
     688            0 :                                     query: format!(
     689            0 :                                         "ALTER ROLE {} {}",
     690            0 :                                         role.name.pg_quote(),
     691            0 :                                         role.to_pg_options(),
     692            0 :                                     ),
     693            0 :                                     comment: None,
     694            0 :                                 })
     695              :                             } else {
     696            0 :                                 None
     697              :                             }
     698              :                         }
     699              :                         None => {
     700            0 :                             let query = if !jwks_roles.contains(role.name.as_str()) {
     701            0 :                                 format!(
     702            0 :                                     "CREATE ROLE {} INHERIT CREATEROLE CREATEDB BYPASSRLS REPLICATION IN ROLE neon_superuser {}",
     703            0 :                                     role.name.pg_quote(),
     704            0 :                                     role.to_pg_options(),
     705            0 :                                 )
     706              :                             } else {
     707            0 :                                 format!(
     708            0 :                                     "CREATE ROLE {} {}",
     709            0 :                                     role.name.pg_quote(),
     710            0 :                                     role.to_pg_options(),
     711            0 :                                 )
     712              :                             };
     713            0 :                             Some(Operation {
     714            0 :                                 query,
     715            0 :                                 comment: Some(format!("creating role {}", role.name)),
     716            0 :                             })
     717              :                         }
     718              :                     }
     719            0 :                 });
     720            0 : 
     721            0 :             Ok(Box::new(operations))
     722              :         }
     723              :         ApplySpecPhase::RenameAndDeleteDatabases => {
     724            0 :             let mut ctx = ctx.write().await;
     725              : 
     726            0 :             let operations = spec
     727            0 :                 .delta_operations
     728            0 :                 .iter()
     729            0 :                 .flatten()
     730            0 :                 .filter_map(move |op| {
     731            0 :                     let databases = &mut ctx.dbs;
     732            0 :                     match op.action.as_str() {
     733            0 :                         // We do not check whether the DB exists or not,
     734            0 :                         // Postgres will take care of it for us
     735            0 :                         "delete_db" => {
     736            0 :                             let (db_name, outer_tag) = op.name.pg_quote_dollar();
     737            0 :                             // In Postgres we can't drop a database if it is a template.
     738            0 :                             // So we need to unset the template flag first, but it could
     739            0 :                             // be a retry, so we could've already dropped the database.
     740            0 :                             // Check that database exists first to make it idempotent.
     741            0 :                             let unset_template_query: String = format!(
     742            0 :                                 include_str!("sql/unset_template_for_drop_dbs.sql"),
     743            0 :                                 datname = db_name,
     744            0 :                                 outer_tag = outer_tag,
     745            0 :                             );
     746            0 : 
     747            0 :                             // Use FORCE to drop database even if there are active connections.
     748            0 :                             // We run this from `cloud_admin`, so it should have enough privileges.
     749            0 :                             //
     750            0 :                             // NB: there could be other db states, which prevent us from dropping
     751            0 :                             // the database. For example, if db is used by any active subscription
     752            0 :                             // or replication slot.
     753            0 :                             // Such cases are handled in the DropLogicalSubscriptions
     754            0 :                             // phase. We do all the cleanup before actually dropping the database.
     755            0 :                             let drop_db_query: String = format!(
     756            0 :                                 "DROP DATABASE IF EXISTS {} WITH (FORCE)",
     757            0 :                                 &op.name.pg_quote()
     758            0 :                             );
     759            0 : 
     760            0 :                             databases.remove(&op.name);
     761            0 : 
     762            0 :                             Some(vec![
     763            0 :                                 Operation {
     764            0 :                                     query: unset_template_query,
     765            0 :                                     comment: Some(format!(
     766            0 :                                         "optionally clearing template flags for DB {}",
     767            0 :                                         op.name,
     768            0 :                                     )),
     769            0 :                                 },
     770            0 :                                 Operation {
     771            0 :                                     query: drop_db_query,
     772            0 :                                     comment: Some(format!("deleting database {}", op.name,)),
     773            0 :                                 },
     774            0 :                             ])
     775              :                         }
     776            0 :                         "rename_db" => {
     777            0 :                             if let Some(mut db) = databases.remove(&op.name) {
     778              :                                 // update state of known databases
     779            0 :                                 let new_name = op.new_name.as_ref().unwrap();
     780            0 :                                 db.name = new_name.clone();
     781            0 :                                 databases.insert(db.name.clone(), db);
     782            0 : 
     783            0 :                                 Some(vec![Operation {
     784            0 :                                     query: format!(
     785            0 :                                         "ALTER DATABASE {} RENAME TO {}",
     786            0 :                                         op.name.pg_quote(),
     787            0 :                                         new_name.pg_quote(),
     788            0 :                                     ),
     789            0 :                                     comment: Some(format!(
     790            0 :                                         "renaming database '{}' to '{}'",
     791            0 :                                         op.name, new_name
     792            0 :                                     )),
     793            0 :                                 }])
     794              :                             } else {
     795            0 :                                 None
     796              :                             }
     797              :                         }
     798            0 :                         _ => None,
     799              :                     }
     800            0 :                 })
     801            0 :                 .flatten();
     802            0 : 
     803            0 :             Ok(Box::new(operations))
     804              :         }
     805              :         ApplySpecPhase::CreateAndAlterDatabases => {
     806            0 :             let mut ctx = ctx.write().await;
     807              : 
     808            0 :             let operations = spec
     809            0 :                 .cluster
     810            0 :                 .databases
     811            0 :                 .iter()
     812            0 :                 .filter_map(move |db| {
     813            0 :                     let databases = &mut ctx.dbs;
     814            0 :                     if let Some(edb) = databases.get_mut(&db.name) {
     815            0 :                         let change_owner = if edb.owner.starts_with('"') {
     816            0 :                             db.owner.pg_quote() != edb.owner
     817              :                         } else {
     818            0 :                             db.owner != edb.owner
     819              :                         };
     820              : 
     821            0 :                         edb.owner = db.owner.clone();
     822            0 : 
     823            0 :                         if change_owner {
     824            0 :                             Some(vec![Operation {
     825            0 :                                 query: format!(
     826            0 :                                     "ALTER DATABASE {} OWNER TO {}",
     827            0 :                                     db.name.pg_quote(),
     828            0 :                                     db.owner.pg_quote()
     829            0 :                                 ),
     830            0 :                                 comment: Some(format!(
     831            0 :                                     "changing database owner of database {} to {}",
     832            0 :                                     db.name, db.owner
     833            0 :                                 )),
     834            0 :                             }])
     835              :                         } else {
     836            0 :                             None
     837              :                         }
     838              :                     } else {
     839            0 :                         databases.insert(db.name.clone(), db.clone());
     840            0 : 
     841            0 :                         Some(vec![
     842            0 :                             Operation {
     843            0 :                                 query: format!(
     844            0 :                                     "CREATE DATABASE {} {}",
     845            0 :                                     db.name.pg_quote(),
     846            0 :                                     db.to_pg_options(),
     847            0 :                                 ),
     848            0 :                                 comment: None,
     849            0 :                             },
     850            0 :                             Operation {
     851            0 :                                 // ALL PRIVILEGES grants CREATE, CONNECT, and TEMPORARY on the database
     852            0 :                                 // (see https://www.postgresql.org/docs/current/ddl-priv.html)
     853            0 :                                 query: format!(
     854            0 :                                     "GRANT ALL PRIVILEGES ON DATABASE {} TO neon_superuser",
     855            0 :                                     db.name.pg_quote()
     856            0 :                                 ),
     857            0 :                                 comment: None,
     858            0 :                             },
     859            0 :                         ])
     860              :                     }
     861            0 :                 })
     862            0 :                 .flatten();
     863            0 : 
     864            0 :             Ok(Box::new(operations))
     865              :         }
     866            0 :         ApplySpecPhase::CreateSchemaNeon => Ok(Box::new(once(Operation {
     867            0 :             query: String::from("CREATE SCHEMA IF NOT EXISTS neon"),
     868            0 :             comment: Some(String::from(
     869            0 :                 "create schema for neon extension and utils tables",
     870            0 :             )),
     871            0 :         }))),
     872            0 :         ApplySpecPhase::RunInEachDatabase { db, subphase } => {
     873              :             // Do some checks that user DB exists and we can access it.
     874              :             //
     875              :             // During the phases like DropLogicalSubscriptions, DeleteDBRoleReferences,
     876              :             // which happen before dropping the DB, the current run could be a retry,
     877              :             // so it's a valid case when DB is absent already. The case of
     878              :             // `pg_database.datallowconn = false`/`restrict_conn` is a bit tricky, as
     879              :             // in theory user can have some dangling objects there, so we will fail at
     880              :             // the actual drop later. Yet, to fix that in the current code we would need
     881              :             // to ALTER DATABASE, and then check back, but that even more invasive, so
     882              :             // that's not what we really want to do here.
     883              :             //
     884              :             // For ChangeSchemaPerms, skipping DBs we cannot access is totally fine.
     885            0 :             if let DB::UserDB(db) = db {
     886            0 :                 let databases = &ctx.read().await.dbs;
     887              : 
     888            0 :                 let edb = match databases.get(&db.name) {
     889            0 :                     Some(edb) => edb,
     890              :                     None => {
     891            0 :                         warn!(
     892            0 :                             "skipping RunInEachDatabase phase {:?}, database {} doesn't exist in PostgreSQL",
     893              :                             subphase, db.name
     894              :                         );
     895            0 :                         return Ok(Box::new(empty()));
     896              :                     }
     897              :                 };
     898              : 
     899            0 :                 if edb.restrict_conn || edb.invalid {
     900            0 :                     warn!(
     901            0 :                         "skipping RunInEachDatabase phase {:?}, database {} is (restrict_conn={}, invalid={})",
     902              :                         subphase, db.name, edb.restrict_conn, edb.invalid
     903              :                     );
     904            0 :                     return Ok(Box::new(empty()));
     905            0 :                 }
     906            0 :             }
     907              : 
     908            0 :             match subphase {
     909              :                 PerDatabasePhase::DropLogicalSubscriptions => {
     910            0 :                     match &db {
     911            0 :                         DB::UserDB(db) => {
     912            0 :                             let (db_name, outer_tag) = db.name.pg_quote_dollar();
     913            0 :                             let drop_subscription_query: String = format!(
     914            0 :                                 include_str!("sql/drop_subscriptions.sql"),
     915            0 :                                 datname_str = db_name,
     916            0 :                                 outer_tag = outer_tag,
     917            0 :                             );
     918            0 : 
     919            0 :                             let operations = vec![Operation {
     920            0 :                                 query: drop_subscription_query,
     921            0 :                                 comment: Some(format!(
     922            0 :                                     "optionally dropping subscriptions for DB {}",
     923            0 :                                     db.name,
     924            0 :                                 )),
     925            0 :                             }]
     926            0 :                             .into_iter();
     927            0 : 
     928            0 :                             Ok(Box::new(operations))
     929              :                         }
     930              :                         // skip this cleanup for the system databases
     931              :                         // because users can't drop them
     932            0 :                         DB::SystemDB => Ok(Box::new(empty())),
     933              :                     }
     934              :                 }
     935              :                 PerDatabasePhase::DeleteDBRoleReferences => {
     936            0 :                     let ctx = ctx.read().await;
     937              : 
     938            0 :                     let operations =
     939            0 :                         spec.delta_operations
     940            0 :                             .iter()
     941            0 :                             .flatten()
     942            0 :                             .filter(|op| op.action == "delete_role")
     943            0 :                             .filter_map(move |op| {
     944            0 :                                 if db.is_owned_by(&op.name) {
     945            0 :                                     return None;
     946            0 :                                 }
     947            0 :                                 if !ctx.roles.contains_key(&op.name) {
     948            0 :                                     return None;
     949            0 :                                 }
     950            0 :                                 let quoted = op.name.pg_quote();
     951            0 :                                 let new_owner = match &db {
     952            0 :                                     DB::SystemDB => PgIdent::from("cloud_admin").pg_quote(),
     953            0 :                                     DB::UserDB(db) => db.owner.pg_quote(),
     954              :                                 };
     955            0 :                                 let (escaped_role, outer_tag) = op.name.pg_quote_dollar();
     956            0 : 
     957            0 :                                 Some(vec![
     958            0 :                                     // This will reassign all dependent objects to the db owner
     959            0 :                                     Operation {
     960            0 :                                         query: format!(
     961            0 :                                             "REASSIGN OWNED BY {} TO {}",
     962            0 :                                             quoted, new_owner,
     963            0 :                                         ),
     964            0 :                                         comment: None,
     965            0 :                                     },
     966            0 :                                     // Revoke some potentially blocking privileges (Neon-specific currently)
     967            0 :                                     Operation {
     968            0 :                                         query: format!(
     969            0 :                                             include_str!("sql/pre_drop_role_revoke_privileges.sql"),
     970            0 :                                             // N.B. this has to be properly dollar-escaped with `pg_quote_dollar()`
     971            0 :                                             role_name = escaped_role,
     972            0 :                                             outer_tag = outer_tag,
     973            0 :                                         ),
     974            0 :                                         comment: None,
     975            0 :                                     },
     976            0 :                                     // This now will only drop privileges of the role
     977            0 :                                     // TODO: this is obviously not 100% true because of the above case,
     978            0 :                                     // there could be still some privileges that are not revoked. Maybe this
     979            0 :                                     // only drops privileges that were granted *by this* role, not *to this* role,
     980            0 :                                     // but this has to be checked.
     981            0 :                                     Operation {
     982            0 :                                         query: format!("DROP OWNED BY {}", quoted),
     983            0 :                                         comment: None,
     984            0 :                                     },
     985            0 :                                 ])
     986            0 :                             })
     987            0 :                             .flatten();
     988            0 : 
     989            0 :                     Ok(Box::new(operations))
     990              :                 }
     991              :                 PerDatabasePhase::ChangeSchemaPerms => {
     992            0 :                     let db = match &db {
     993              :                         // ignore schema permissions on the system database
     994            0 :                         DB::SystemDB => return Ok(Box::new(empty())),
     995            0 :                         DB::UserDB(db) => db,
     996            0 :                     };
     997            0 :                     let (db_owner, outer_tag) = db.owner.pg_quote_dollar();
     998            0 : 
     999            0 :                     let operations = vec![
    1000            0 :                         Operation {
    1001            0 :                             query: format!(
    1002            0 :                                 include_str!("sql/set_public_schema_owner.sql"),
    1003            0 :                                 db_owner = db_owner,
    1004            0 :                                 outer_tag = outer_tag,
    1005            0 :                             ),
    1006            0 :                             comment: None,
    1007            0 :                         },
    1008            0 :                         Operation {
    1009            0 :                             query: String::from(include_str!("sql/default_grants.sql")),
    1010            0 :                             comment: None,
    1011            0 :                         },
    1012            0 :                     ]
    1013            0 :                     .into_iter();
    1014            0 : 
    1015            0 :                     Ok(Box::new(operations))
    1016              :                 }
    1017              :                 // TODO: remove this completely https://github.com/neondatabase/cloud/issues/22663
    1018              :                 PerDatabasePhase::HandleAnonExtension => {
    1019              :                     // Only install Anon into user databases
    1020            0 :                     let db = match &db {
    1021            0 :                         DB::SystemDB => return Ok(Box::new(empty())),
    1022            0 :                         DB::UserDB(db) => db,
    1023            0 :                     };
    1024            0 :                     // Never install Anon when it's not enabled as feature
    1025            0 :                     if !spec.features.contains(&ComputeFeature::AnonExtension) {
    1026            0 :                         return Ok(Box::new(empty()));
    1027            0 :                     }
    1028            0 : 
    1029            0 :                     // Only install Anon when it's added in preload libraries
    1030            0 :                     let opt_libs = spec.cluster.settings.find("shared_preload_libraries");
    1031              : 
    1032            0 :                     let libs = match opt_libs {
    1033            0 :                         Some(libs) => libs,
    1034            0 :                         None => return Ok(Box::new(empty())),
    1035              :                     };
    1036              : 
    1037            0 :                     if !libs.contains("anon") {
    1038            0 :                         return Ok(Box::new(empty()));
    1039            0 :                     }
    1040            0 : 
    1041            0 :                     let db_owner = db.owner.pg_quote();
    1042            0 : 
    1043            0 :                     let operations = vec![
    1044            0 :                         // Create anon extension if this compute needs it
    1045            0 :                         // Users cannot create it themselves, because superuser is required.
    1046            0 :                         Operation {
    1047            0 :                             query: String::from("CREATE EXTENSION IF NOT EXISTS anon CASCADE"),
    1048            0 :                             comment: Some(String::from("creating anon extension")),
    1049            0 :                         },
    1050            0 :                         // Initialize anon extension
    1051            0 :                         // This also requires superuser privileges, so users cannot do it themselves.
    1052            0 :                         Operation {
    1053            0 :                             query: String::from("SELECT anon.init()"),
    1054            0 :                             comment: Some(String::from("initializing anon extension data")),
    1055            0 :                         },
    1056            0 :                         Operation {
    1057            0 :                             query: format!("GRANT ALL ON SCHEMA anon TO {}", db_owner),
    1058            0 :                             comment: Some(String::from(
    1059            0 :                                 "granting anon extension schema permissions",
    1060            0 :                             )),
    1061            0 :                         },
    1062            0 :                         Operation {
    1063            0 :                             query: format!(
    1064            0 :                                 "GRANT ALL ON ALL FUNCTIONS IN SCHEMA anon TO {}",
    1065            0 :                                 db_owner
    1066            0 :                             ),
    1067            0 :                             comment: Some(String::from(
    1068            0 :                                 "granting anon extension schema functions permissions",
    1069            0 :                             )),
    1070            0 :                         },
    1071            0 :                         // We need this, because some functions are defined as SECURITY DEFINER.
    1072            0 :                         // In Postgres SECURITY DEFINER functions are executed with the privileges
    1073            0 :                         // of the owner.
    1074            0 :                         // In anon extension this it is needed to access some GUCs, which are only accessible to
    1075            0 :                         // superuser. But we've patched postgres to allow db_owner to access them as well.
    1076            0 :                         // So we need to change owner of these functions to db_owner.
    1077            0 :                         Operation {
    1078            0 :                             query: format!(
    1079            0 :                                 include_str!("sql/anon_ext_fn_reassign.sql"),
    1080            0 :                                 db_owner = db_owner,
    1081            0 :                             ),
    1082            0 :                             comment: Some(String::from(
    1083            0 :                                 "change anon extension functions owner to database_owner",
    1084            0 :                             )),
    1085            0 :                         },
    1086            0 :                         Operation {
    1087            0 :                             query: format!(
    1088            0 :                                 "GRANT ALL ON ALL TABLES IN SCHEMA anon TO {}",
    1089            0 :                                 db_owner,
    1090            0 :                             ),
    1091            0 :                             comment: Some(String::from(
    1092            0 :                                 "granting anon extension tables permissions",
    1093            0 :                             )),
    1094            0 :                         },
    1095            0 :                         Operation {
    1096            0 :                             query: format!(
    1097            0 :                                 "GRANT ALL ON ALL SEQUENCES IN SCHEMA anon TO {}",
    1098            0 :                                 db_owner,
    1099            0 :                             ),
    1100            0 :                             comment: Some(String::from(
    1101            0 :                                 "granting anon extension sequences permissions",
    1102            0 :                             )),
    1103            0 :                         },
    1104            0 :                     ]
    1105            0 :                     .into_iter();
    1106            0 : 
    1107            0 :                     Ok(Box::new(operations))
    1108              :                 }
    1109              :             }
    1110              :         }
    1111              :         // Interestingly, we only install p_s_s in the main database, even when
    1112              :         // it's preloaded.
    1113              :         ApplySpecPhase::HandleOtherExtensions => {
    1114            0 :             if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
    1115            0 :                 if libs.contains("pg_stat_statements") {
    1116            0 :                     return Ok(Box::new(once(Operation {
    1117            0 :                         query: String::from("CREATE EXTENSION IF NOT EXISTS pg_stat_statements"),
    1118            0 :                         comment: Some(String::from("create system extensions")),
    1119            0 :                     })));
    1120            0 :                 }
    1121            0 :             }
    1122            0 :             Ok(Box::new(empty()))
    1123              :         }
    1124            0 :         ApplySpecPhase::CreatePgauditExtension => Ok(Box::new(once(Operation {
    1125            0 :             query: String::from("CREATE EXTENSION IF NOT EXISTS pgaudit"),
    1126            0 :             comment: Some(String::from("create pgaudit extensions")),
    1127            0 :         }))),
    1128            0 :         ApplySpecPhase::CreatePgauditlogtofileExtension => Ok(Box::new(once(Operation {
    1129            0 :             query: String::from("CREATE EXTENSION IF NOT EXISTS pgauditlogtofile"),
    1130            0 :             comment: Some(String::from("create pgauditlogtofile extensions")),
    1131            0 :         }))),
    1132              :         // Disable pgaudit logging for postgres database.
    1133              :         // Postgres is neon system database used by monitors
    1134              :         // and compute_ctl tuning functions and thus generates a lot of noise.
    1135              :         // We do not consider data stored in this database as sensitive.
    1136              :         ApplySpecPhase::DisablePostgresDBPgAudit => {
    1137            0 :             let query = "ALTER DATABASE postgres SET pgaudit.log to 'none'";
    1138            0 :             Ok(Box::new(once(Operation {
    1139            0 :                 query: query.to_string(),
    1140            0 :                 comment: Some(query.to_string()),
    1141            0 :             })))
    1142              :         }
    1143              :         ApplySpecPhase::HandleNeonExtension => {
    1144            0 :             let operations = vec![
    1145            0 :                 Operation {
    1146            0 :                     query: String::from("CREATE EXTENSION IF NOT EXISTS neon WITH SCHEMA neon"),
    1147            0 :                     comment: Some(String::from(
    1148            0 :                         "init: install the extension if not already installed",
    1149            0 :                     )),
    1150            0 :                 },
    1151            0 :                 Operation {
    1152            0 :                     query: String::from(
    1153            0 :                         "UPDATE pg_extension SET extrelocatable = true WHERE extname = 'neon'",
    1154            0 :                     ),
    1155            0 :                     comment: Some(String::from("compat/fix: make neon relocatable")),
    1156            0 :                 },
    1157            0 :                 Operation {
    1158            0 :                     query: String::from("ALTER EXTENSION neon SET SCHEMA neon"),
    1159            0 :                     comment: Some(String::from("compat/fix: alter neon extension schema")),
    1160            0 :                 },
    1161            0 :                 Operation {
    1162            0 :                     query: String::from("ALTER EXTENSION neon UPDATE"),
    1163            0 :                     comment: Some(String::from("compat/update: update neon extension version")),
    1164            0 :                 },
    1165            0 :             ]
    1166            0 :             .into_iter();
    1167            0 : 
    1168            0 :             Ok(Box::new(operations))
    1169              :         }
    1170            0 :         ApplySpecPhase::CreateAvailabilityCheck => Ok(Box::new(once(Operation {
    1171            0 :             query: String::from(include_str!("sql/add_availabilitycheck_tables.sql")),
    1172            0 :             comment: None,
    1173            0 :         }))),
    1174              :         ApplySpecPhase::DropRoles => {
    1175            0 :             let operations = spec
    1176            0 :                 .delta_operations
    1177            0 :                 .iter()
    1178            0 :                 .flatten()
    1179            0 :                 .filter(|op| op.action == "delete_role")
    1180            0 :                 .map(|op| Operation {
    1181            0 :                     query: format!("DROP ROLE IF EXISTS {}", op.name.pg_quote()),
    1182            0 :                     comment: None,
    1183            0 :                 });
    1184            0 : 
    1185            0 :             Ok(Box::new(operations))
    1186              :         }
    1187            0 :         ApplySpecPhase::FinalizeDropLogicalSubscriptions => Ok(Box::new(once(Operation {
    1188            0 :             query: String::from(include_str!("sql/finalize_drop_subscriptions.sql")),
    1189            0 :             comment: None,
    1190            0 :         }))),
    1191              :     }
    1192            0 : }
        

Generated by: LCOV version 2.1-beta