LCOV - code coverage report
Current view: top level - compute_tools/src - spec.rs (source / functions) Coverage Total Hit
Test: 727bdccc1d7d53837da843959afb612f56da4e79.info Lines: 0.0 % 101 0
Test Date: 2025-01-30 15:18:43 Functions: 0.0 % 5 0

            Line data    Source code
       1              : use anyhow::{anyhow, bail, Result};
       2              : use postgres::Client;
       3              : use reqwest::StatusCode;
       4              : use std::fs::File;
       5              : use std::path::Path;
       6              : use tracing::{error, info, instrument, warn};
       7              : 
       8              : use crate::config;
       9              : use crate::metrics::{CPlaneRequestRPC, CPLANE_REQUESTS_TOTAL, UNKNOWN_HTTP_STATUS};
      10              : use crate::migration::MigrationRunner;
      11              : use crate::params::PG_HBA_ALL_MD5;
      12              : use crate::pg_helpers::*;
      13              : 
      14              : use compute_api::responses::{ControlPlaneComputeStatus, ControlPlaneSpecResponse};
      15              : use compute_api::spec::ComputeSpec;
      16              : 
      17              : // Do control plane request and return response if any. In case of error it
      18              : // returns a bool flag indicating whether it makes sense to retry the request
      19              : // and a string with error message.
      20            0 : fn do_control_plane_request(
      21            0 :     uri: &str,
      22            0 :     jwt: &str,
      23            0 : ) -> Result<ControlPlaneSpecResponse, (bool, String, String)> {
      24            0 :     let resp = reqwest::blocking::Client::new()
      25            0 :         .get(uri)
      26            0 :         .header("Authorization", format!("Bearer {}", jwt))
      27            0 :         .send()
      28            0 :         .map_err(|e| {
      29            0 :             (
      30            0 :                 true,
      31            0 :                 format!("could not perform spec request to control plane: {}", e),
      32            0 :                 UNKNOWN_HTTP_STATUS.to_string(),
      33            0 :             )
      34            0 :         })?;
      35              : 
      36            0 :     let status = resp.status();
      37            0 :     match status {
      38            0 :         StatusCode::OK => match resp.json::<ControlPlaneSpecResponse>() {
      39            0 :             Ok(spec_resp) => Ok(spec_resp),
      40            0 :             Err(e) => Err((
      41            0 :                 true,
      42            0 :                 format!("could not deserialize control plane response: {}", e),
      43            0 :                 status.to_string(),
      44            0 :             )),
      45              :         },
      46            0 :         StatusCode::SERVICE_UNAVAILABLE => Err((
      47            0 :             true,
      48            0 :             "control plane is temporarily unavailable".to_string(),
      49            0 :             status.to_string(),
      50            0 :         )),
      51              :         StatusCode::BAD_GATEWAY => {
      52              :             // We have a problem with intermittent 502 errors now
      53              :             // https://github.com/neondatabase/cloud/issues/2353
      54              :             // It's fine to retry GET request in this case.
      55            0 :             Err((
      56            0 :                 true,
      57            0 :                 "control plane request failed with 502".to_string(),
      58            0 :                 status.to_string(),
      59            0 :             ))
      60              :         }
      61              :         // Another code, likely 500 or 404, means that compute is unknown to the control plane
      62              :         // or some internal failure happened. Doesn't make much sense to retry in this case.
      63            0 :         _ => Err((
      64            0 :             false,
      65            0 :             format!("unexpected control plane response status code: {}", status),
      66            0 :             status.to_string(),
      67            0 :         )),
      68              :     }
      69            0 : }
      70              : 
      71              : /// Request spec from the control-plane by compute_id. If `NEON_CONTROL_PLANE_TOKEN`
      72              : /// env variable is set, it will be used for authorization.
      73            0 : pub fn get_spec_from_control_plane(
      74            0 :     base_uri: &str,
      75            0 :     compute_id: &str,
      76            0 : ) -> Result<Option<ComputeSpec>> {
      77            0 :     let cp_uri = format!("{base_uri}/compute/api/v2/computes/{compute_id}/spec");
      78            0 :     let jwt: String = match std::env::var("NEON_CONTROL_PLANE_TOKEN") {
      79            0 :         Ok(v) => v,
      80            0 :         Err(_) => "".to_string(),
      81              :     };
      82            0 :     let mut attempt = 1;
      83            0 :     let mut spec: Result<Option<ComputeSpec>> = Ok(None);
      84            0 : 
      85            0 :     info!("getting spec from control plane: {}", cp_uri);
      86              : 
      87              :     // Do 3 attempts to get spec from the control plane using the following logic:
      88              :     // - network error -> then retry
      89              :     // - compute id is unknown or any other error -> bail out
      90              :     // - no spec for compute yet (Empty state) -> return Ok(None)
      91              :     // - got spec -> return Ok(Some(spec))
      92            0 :     while attempt < 4 {
      93            0 :         spec = match do_control_plane_request(&cp_uri, &jwt) {
      94            0 :             Ok(spec_resp) => {
      95            0 :                 CPLANE_REQUESTS_TOTAL
      96            0 :                     .with_label_values(&[
      97            0 :                         CPlaneRequestRPC::GetSpec.as_str(),
      98            0 :                         &StatusCode::OK.to_string(),
      99            0 :                     ])
     100            0 :                     .inc();
     101            0 :                 match spec_resp.status {
     102            0 :                     ControlPlaneComputeStatus::Empty => Ok(None),
     103              :                     ControlPlaneComputeStatus::Attached => {
     104            0 :                         if let Some(spec) = spec_resp.spec {
     105            0 :                             Ok(Some(spec))
     106              :                         } else {
     107            0 :                             bail!("compute is attached, but spec is empty")
     108              :                         }
     109              :                     }
     110              :                 }
     111              :             }
     112            0 :             Err((retry, msg, status)) => {
     113            0 :                 CPLANE_REQUESTS_TOTAL
     114            0 :                     .with_label_values(&[CPlaneRequestRPC::GetSpec.as_str(), &status])
     115            0 :                     .inc();
     116            0 :                 if retry {
     117            0 :                     Err(anyhow!(msg))
     118              :                 } else {
     119            0 :                     bail!(msg);
     120              :                 }
     121              :             }
     122              :         };
     123              : 
     124            0 :         if let Err(e) = &spec {
     125            0 :             error!("attempt {} to get spec failed with: {}", attempt, e);
     126              :         } else {
     127            0 :             return spec;
     128              :         }
     129              : 
     130            0 :         attempt += 1;
     131            0 :         std::thread::sleep(std::time::Duration::from_millis(100));
     132              :     }
     133              : 
     134              :     // All attempts failed, return error.
     135            0 :     spec
     136            0 : }
     137              : 
     138              : /// Check `pg_hba.conf` and update if needed to allow external connections.
     139            0 : pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
     140            0 :     // XXX: consider making it a part of spec.json
     141            0 :     info!("checking pg_hba.conf");
     142            0 :     let pghba_path = pgdata_path.join("pg_hba.conf");
     143            0 : 
     144            0 :     if config::line_in_file(&pghba_path, PG_HBA_ALL_MD5)? {
     145            0 :         info!("updated pg_hba.conf to allow external connections");
     146              :     } else {
     147            0 :         info!("pg_hba.conf is up-to-date");
     148              :     }
     149              : 
     150            0 :     Ok(())
     151            0 : }
     152              : 
     153              : /// Create a standby.signal file
     154            0 : pub fn add_standby_signal(pgdata_path: &Path) -> Result<()> {
     155            0 :     // XXX: consider making it a part of spec.json
     156            0 :     info!("adding standby.signal");
     157            0 :     let signalfile = pgdata_path.join("standby.signal");
     158            0 : 
     159            0 :     if !signalfile.exists() {
     160            0 :         info!("created standby.signal");
     161            0 :         File::create(signalfile)?;
     162              :     } else {
     163            0 :         info!("reused pre-existing standby.signal");
     164              :     }
     165            0 :     Ok(())
     166            0 : }
     167              : 
     168              : #[instrument(skip_all)]
     169              : pub fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
     170              :     info!("handle neon extension upgrade");
     171              :     let query = "ALTER EXTENSION neon UPDATE";
     172              :     info!("update neon extension version with query: {}", query);
     173              :     client.simple_query(query)?;
     174              : 
     175              :     Ok(())
     176              : }
     177              : 
     178              : #[instrument(skip_all)]
     179              : pub fn handle_migrations(client: &mut Client) -> Result<()> {
     180              :     info!("handle migrations");
     181              : 
     182              :     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
     183              :     // !BE SURE TO ONLY ADD MIGRATIONS TO THE END OF THIS ARRAY. IF YOU DO NOT, VERY VERY BAD THINGS MAY HAPPEN!
     184              :     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
     185              : 
     186              :     // Add new migrations in numerical order.
     187              :     let migrations = [
     188              :         include_str!("./migrations/0001-neon_superuser_bypass_rls.sql"),
     189              :         include_str!("./migrations/0002-alter_roles.sql"),
     190              :         include_str!("./migrations/0003-grant_pg_create_subscription_to_neon_superuser.sql"),
     191              :         include_str!("./migrations/0004-grant_pg_monitor_to_neon_superuser.sql"),
     192              :         include_str!("./migrations/0005-grant_all_on_tables_to_neon_superuser.sql"),
     193              :         include_str!("./migrations/0006-grant_all_on_sequences_to_neon_superuser.sql"),
     194              :         include_str!(
     195              :             "./migrations/0007-grant_all_on_tables_to_neon_superuser_with_grant_option.sql"
     196              :         ),
     197              :         include_str!(
     198              :             "./migrations/0008-grant_all_on_sequences_to_neon_superuser_with_grant_option.sql"
     199              :         ),
     200              :         include_str!("./migrations/0009-revoke_replication_for_previously_allowed_roles.sql"),
     201              :         include_str!(
     202              :             "./migrations/0010-grant_snapshot_synchronization_funcs_to_neon_superuser.sql"
     203              :         ),
     204              :         include_str!(
     205              :             "./migrations/0011-grant_pg_show_replication_origin_status_to_neon_superuser.sql"
     206              :         ),
     207              :     ];
     208              : 
     209              :     MigrationRunner::new(client, &migrations).run_migrations()?;
     210              : 
     211              :     Ok(())
     212              : }
     213              : 
     214              : /// Connect to the database as superuser and pre-create anon extension
     215              : /// if it is present in shared_preload_libraries
     216              : #[instrument(skip_all)]
     217              : pub fn handle_extension_anon(
     218              :     spec: &ComputeSpec,
     219              :     db_owner: &str,
     220              :     db_client: &mut Client,
     221              :     grants_only: bool,
     222              : ) -> Result<()> {
     223              :     info!("handle extension anon");
     224              : 
     225              :     if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
     226              :         if libs.contains("anon") {
     227              :             if !grants_only {
     228              :                 // check if extension is already initialized using anon.is_initialized()
     229              :                 let query = "SELECT anon.is_initialized()";
     230              :                 match db_client.query(query, &[]) {
     231              :                     Ok(rows) => {
     232              :                         if !rows.is_empty() {
     233              :                             let is_initialized: bool = rows[0].get(0);
     234              :                             if is_initialized {
     235              :                                 info!("anon extension is already initialized");
     236              :                                 return Ok(());
     237              :                             }
     238              :                         }
     239              :                     }
     240              :                     Err(e) => {
     241              :                         warn!(
     242              :                             "anon extension is_installed check failed with expected error: {}",
     243              :                             e
     244              :                         );
     245              :                     }
     246              :                 };
     247              : 
     248              :                 // Create anon extension if this compute needs it
     249              :                 // Users cannot create it themselves, because superuser is required.
     250              :                 let mut query = "CREATE EXTENSION IF NOT EXISTS anon CASCADE";
     251              :                 info!("creating anon extension with query: {}", query);
     252              :                 match db_client.query(query, &[]) {
     253              :                     Ok(_) => {}
     254              :                     Err(e) => {
     255              :                         error!("anon extension creation failed with error: {}", e);
     256              :                         return Ok(());
     257              :                     }
     258              :                 }
     259              : 
     260              :                 // check that extension is installed
     261              :                 query = "SELECT extname FROM pg_extension WHERE extname = 'anon'";
     262              :                 let rows = db_client.query(query, &[])?;
     263              :                 if rows.is_empty() {
     264              :                     error!("anon extension is not installed");
     265              :                     return Ok(());
     266              :                 }
     267              : 
     268              :                 // Initialize anon extension
     269              :                 // This also requires superuser privileges, so users cannot do it themselves.
     270              :                 query = "SELECT anon.init()";
     271              :                 match db_client.query(query, &[]) {
     272              :                     Ok(_) => {}
     273              :                     Err(e) => {
     274              :                         error!("anon.init() failed with error: {}", e);
     275              :                         return Ok(());
     276              :                     }
     277              :                 }
     278              :             }
     279              : 
     280              :             // check that extension is installed, if not bail early
     281              :             let query = "SELECT extname FROM pg_extension WHERE extname = 'anon'";
     282              :             match db_client.query(query, &[]) {
     283              :                 Ok(rows) => {
     284              :                     if rows.is_empty() {
     285              :                         error!("anon extension is not installed");
     286              :                         return Ok(());
     287              :                     }
     288              :                 }
     289              :                 Err(e) => {
     290              :                     error!("anon extension check failed with error: {}", e);
     291              :                     return Ok(());
     292              :                 }
     293              :             };
     294              : 
     295              :             let query = format!("GRANT ALL ON SCHEMA anon TO {}", db_owner);
     296              :             info!("granting anon extension permissions with query: {}", query);
     297              :             db_client.simple_query(&query)?;
     298              : 
     299              :             // Grant permissions to db_owner to use anon extension functions
     300              :             let query = format!("GRANT ALL ON ALL FUNCTIONS IN SCHEMA anon TO {}", db_owner);
     301              :             info!("granting anon extension permissions with query: {}", query);
     302              :             db_client.simple_query(&query)?;
     303              : 
     304              :             // This is needed, because some functions are defined as SECURITY DEFINER.
     305              :             // In Postgres SECURITY DEFINER functions are executed with the privileges
     306              :             // of the owner.
     307              :             // In anon extension this it is needed to access some GUCs, which are only accessible to
     308              :             // superuser. But we've patched postgres to allow db_owner to access them as well.
     309              :             // So we need to change owner of these functions to db_owner.
     310              :             let query = format!("
     311              :                 SELECT 'ALTER FUNCTION '||nsp.nspname||'.'||p.proname||'('||pg_get_function_identity_arguments(p.oid)||') OWNER TO {};'
     312              :                 from pg_proc p
     313              :                 join pg_namespace nsp ON p.pronamespace = nsp.oid
     314              :                 where nsp.nspname = 'anon';", db_owner);
     315              : 
     316              :             info!("change anon extension functions owner to db owner");
     317              :             db_client.simple_query(&query)?;
     318              : 
     319              :             //  affects views as well
     320              :             let query = format!("GRANT ALL ON ALL TABLES IN SCHEMA anon TO {}", db_owner);
     321              :             info!("granting anon extension permissions with query: {}", query);
     322              :             db_client.simple_query(&query)?;
     323              : 
     324              :             let query = format!("GRANT ALL ON ALL SEQUENCES IN SCHEMA anon TO {}", db_owner);
     325              :             info!("granting anon extension permissions with query: {}", query);
     326              :             db_client.simple_query(&query)?;
     327              :         }
     328              :     }
     329              : 
     330              :     Ok(())
     331              : }
        

Generated by: LCOV version 2.1-beta