LCOV - code coverage report
Current view: top level - compute_tools/src - spec.rs (source / functions) Coverage Total Hit
Test: 5e392a02abbad1ab595f4dba672e219a49f7f539.info Lines: 0.0 % 95 0
Test Date: 2025-04-11 22:43:24 Functions: 0.0 % 5 0

            Line data    Source code
       1              : use std::fs::File;
       2              : use std::path::Path;
       3              : 
       4              : use anyhow::{Result, anyhow, bail};
       5              : use compute_api::responses::{
       6              :     ComputeConfig, ControlPlaneComputeStatus, ControlPlaneConfigResponse,
       7              : };
       8              : use reqwest::StatusCode;
       9              : use tokio_postgres::Client;
      10              : use tracing::{error, info, instrument};
      11              : 
      12              : use crate::config;
      13              : use crate::metrics::{CPLANE_REQUESTS_TOTAL, CPlaneRequestRPC, UNKNOWN_HTTP_STATUS};
      14              : use crate::migration::MigrationRunner;
      15              : use crate::params::PG_HBA_ALL_MD5;
      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<ControlPlaneConfigResponse, (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 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::<ControlPlaneConfigResponse>() {
      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 config from the control-plane by compute_id. If
      72              : /// `NEON_CONTROL_PLANE_TOKEN` env variable is set, it will be used for
      73              : /// authorization.
      74            0 : pub fn get_config_from_control_plane(base_uri: &str, compute_id: &str) -> Result<ComputeConfig> {
      75            0 :     let cp_uri = format!("{base_uri}/compute/api/v2/computes/{compute_id}/spec");
      76            0 :     let jwt: String = std::env::var("NEON_CONTROL_PLANE_TOKEN").unwrap_or_default();
      77            0 :     let mut attempt = 1;
      78            0 : 
      79            0 :     info!("getting config from control plane: {}", cp_uri);
      80              : 
      81              :     // Do 3 attempts to get spec from the control plane using the following logic:
      82              :     // - network error -> then retry
      83              :     // - compute id is unknown or any other error -> bail out
      84              :     // - no spec for compute yet (Empty state) -> return Ok(None)
      85              :     // - got config -> return Ok(Some(config))
      86            0 :     while attempt < 4 {
      87            0 :         let result = match do_control_plane_request(&cp_uri, &jwt) {
      88            0 :             Ok(config_resp) => {
      89            0 :                 CPLANE_REQUESTS_TOTAL
      90            0 :                     .with_label_values(&[
      91            0 :                         CPlaneRequestRPC::GetConfig.as_str(),
      92            0 :                         &StatusCode::OK.to_string(),
      93            0 :                     ])
      94            0 :                     .inc();
      95            0 :                 match config_resp.status {
      96            0 :                     ControlPlaneComputeStatus::Empty => Ok(config_resp.into()),
      97              :                     ControlPlaneComputeStatus::Attached => {
      98            0 :                         if config_resp.spec.is_some() {
      99            0 :                             Ok(config_resp.into())
     100              :                         } else {
     101            0 :                             bail!("compute is attached, but spec is empty")
     102              :                         }
     103              :                     }
     104              :                 }
     105              :             }
     106            0 :             Err((retry, msg, status)) => {
     107            0 :                 CPLANE_REQUESTS_TOTAL
     108            0 :                     .with_label_values(&[CPlaneRequestRPC::GetConfig.as_str(), &status])
     109            0 :                     .inc();
     110            0 :                 if retry {
     111            0 :                     Err(anyhow!(msg))
     112              :                 } else {
     113            0 :                     bail!(msg);
     114              :                 }
     115              :             }
     116              :         };
     117              : 
     118            0 :         if let Err(e) = &result {
     119            0 :             error!("attempt {} to get config failed with: {}", attempt, e);
     120              :         } else {
     121            0 :             return result;
     122              :         }
     123              : 
     124            0 :         attempt += 1;
     125            0 :         std::thread::sleep(std::time::Duration::from_millis(100));
     126              :     }
     127              : 
     128              :     // All attempts failed, return error.
     129            0 :     Err(anyhow::anyhow!(
     130            0 :         "Exhausted all attempts to retrieve the config from the control plane"
     131            0 :     ))
     132            0 : }
     133              : 
     134              : /// Check `pg_hba.conf` and update if needed to allow external connections.
     135            0 : pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
     136            0 :     // XXX: consider making it a part of config.json
     137            0 :     let pghba_path = pgdata_path.join("pg_hba.conf");
     138            0 : 
     139            0 :     if config::line_in_file(&pghba_path, PG_HBA_ALL_MD5)? {
     140            0 :         info!("updated pg_hba.conf to allow external connections");
     141              :     } else {
     142            0 :         info!("pg_hba.conf is up-to-date");
     143              :     }
     144              : 
     145            0 :     Ok(())
     146            0 : }
     147              : 
     148              : /// Create a standby.signal file
     149            0 : pub fn add_standby_signal(pgdata_path: &Path) -> Result<()> {
     150            0 :     // XXX: consider making it a part of config.json
     151            0 :     let signalfile = pgdata_path.join("standby.signal");
     152            0 : 
     153            0 :     if !signalfile.exists() {
     154            0 :         File::create(signalfile)?;
     155            0 :         info!("created standby.signal");
     156              :     } else {
     157            0 :         info!("reused pre-existing standby.signal");
     158              :     }
     159            0 :     Ok(())
     160            0 : }
     161              : 
     162              : #[instrument(skip_all)]
     163              : pub async fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
     164              :     let query = "ALTER EXTENSION neon UPDATE";
     165              :     info!("update neon extension version with query: {}", query);
     166              :     client.simple_query(query).await?;
     167              : 
     168              :     Ok(())
     169              : }
     170              : 
     171              : #[instrument(skip_all)]
     172              : pub async fn handle_migrations(client: &mut Client) -> Result<()> {
     173              :     info!("handle migrations");
     174              : 
     175              :     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
     176              :     // !BE SURE TO ONLY ADD MIGRATIONS TO THE END OF THIS ARRAY. IF YOU DO NOT, VERY VERY BAD THINGS MAY HAPPEN!
     177              :     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
     178              : 
     179              :     // Add new migrations in numerical order.
     180              :     let migrations = [
     181              :         include_str!("./migrations/0001-neon_superuser_bypass_rls.sql"),
     182              :         include_str!("./migrations/0002-alter_roles.sql"),
     183              :         include_str!("./migrations/0003-grant_pg_create_subscription_to_neon_superuser.sql"),
     184              :         include_str!("./migrations/0004-grant_pg_monitor_to_neon_superuser.sql"),
     185              :         include_str!("./migrations/0005-grant_all_on_tables_to_neon_superuser.sql"),
     186              :         include_str!("./migrations/0006-grant_all_on_sequences_to_neon_superuser.sql"),
     187              :         include_str!(
     188              :             "./migrations/0007-grant_all_on_tables_to_neon_superuser_with_grant_option.sql"
     189              :         ),
     190              :         include_str!(
     191              :             "./migrations/0008-grant_all_on_sequences_to_neon_superuser_with_grant_option.sql"
     192              :         ),
     193              :         include_str!("./migrations/0009-revoke_replication_for_previously_allowed_roles.sql"),
     194              :         include_str!(
     195              :             "./migrations/0010-grant_snapshot_synchronization_funcs_to_neon_superuser.sql"
     196              :         ),
     197              :         include_str!(
     198              :             "./migrations/0011-grant_pg_show_replication_origin_status_to_neon_superuser.sql"
     199              :         ),
     200              :     ];
     201              : 
     202              :     MigrationRunner::new(client, &migrations)
     203              :         .run_migrations()
     204              :         .await?;
     205              : 
     206              :     Ok(())
     207              : }
        

Generated by: LCOV version 2.1-beta