LCOV - code coverage report
Current view: top level - compute_tools/src - spec.rs (source / functions) Coverage Total Hit
Test: ae4948feae6a1d420c855050eb8c189119446a71.info Lines: 0.0 % 100 0
Test Date: 2025-03-18 18:33:46 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              :     ComputeCtlConfig, ControlPlaneComputeStatus, ControlPlaneSpecResponse,
       7              : };
       8              : use compute_api::spec::ComputeSpec;
       9              : use reqwest::StatusCode;
      10              : use tokio_postgres::Client;
      11              : use tracing::{error, info, instrument};
      12              : 
      13              : use crate::config;
      14              : use crate::metrics::{CPLANE_REQUESTS_TOTAL, CPlaneRequestRPC, UNKNOWN_HTTP_STATUS};
      15              : use crate::migration::MigrationRunner;
      16              : use crate::params::PG_HBA_ALL_MD5;
      17              : 
      18              : // Do control plane request and return response if any. In case of error it
      19              : // returns a bool flag indicating whether it makes sense to retry the request
      20              : // and a string with error message.
      21            0 : fn do_control_plane_request(
      22            0 :     uri: &str,
      23            0 :     jwt: &str,
      24            0 : ) -> Result<ControlPlaneSpecResponse, (bool, String, String)> {
      25            0 :     let resp = reqwest::blocking::Client::new()
      26            0 :         .get(uri)
      27            0 :         .header("Authorization", format!("Bearer {}", jwt))
      28            0 :         .send()
      29            0 :         .map_err(|e| {
      30            0 :             (
      31            0 :                 true,
      32            0 :                 format!("could not perform spec request to control plane: {:?}", e),
      33            0 :                 UNKNOWN_HTTP_STATUS.to_string(),
      34            0 :             )
      35            0 :         })?;
      36              : 
      37            0 :     let status = resp.status();
      38            0 :     match status {
      39            0 :         StatusCode::OK => match resp.json::<ControlPlaneSpecResponse>() {
      40            0 :             Ok(spec_resp) => Ok(spec_resp),
      41            0 :             Err(e) => Err((
      42            0 :                 true,
      43            0 :                 format!("could not deserialize control plane response: {:?}", e),
      44            0 :                 status.to_string(),
      45            0 :             )),
      46              :         },
      47            0 :         StatusCode::SERVICE_UNAVAILABLE => Err((
      48            0 :             true,
      49            0 :             "control plane is temporarily unavailable".to_string(),
      50            0 :             status.to_string(),
      51            0 :         )),
      52              :         StatusCode::BAD_GATEWAY => {
      53              :             // We have a problem with intermittent 502 errors now
      54              :             // https://github.com/neondatabase/cloud/issues/2353
      55              :             // It's fine to retry GET request in this case.
      56            0 :             Err((
      57            0 :                 true,
      58            0 :                 "control plane request failed with 502".to_string(),
      59            0 :                 status.to_string(),
      60            0 :             ))
      61              :         }
      62              :         // Another code, likely 500 or 404, means that compute is unknown to the control plane
      63              :         // or some internal failure happened. Doesn't make much sense to retry in this case.
      64            0 :         _ => Err((
      65            0 :             false,
      66            0 :             format!("unexpected control plane response status code: {}", status),
      67            0 :             status.to_string(),
      68            0 :         )),
      69              :     }
      70            0 : }
      71              : 
      72              : /// Request spec from the control-plane by compute_id. If `NEON_CONTROL_PLANE_TOKEN`
      73              : /// env variable is set, it will be used for authorization.
      74            0 : pub fn get_spec_from_control_plane(
      75            0 :     base_uri: &str,
      76            0 :     compute_id: &str,
      77            0 : ) -> Result<(Option<ComputeSpec>, ComputeCtlConfig)> {
      78            0 :     let cp_uri = format!("{base_uri}/compute/api/v2/computes/{compute_id}/spec");
      79            0 :     let jwt: String = match std::env::var("NEON_CONTROL_PLANE_TOKEN") {
      80            0 :         Ok(v) => v,
      81            0 :         Err(_) => "".to_string(),
      82              :     };
      83            0 :     let mut attempt = 1;
      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 :         let result = 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, spec_resp.compute_ctl_config)),
     103              :                     ControlPlaneComputeStatus::Attached => {
     104            0 :                         if let Some(spec) = spec_resp.spec {
     105            0 :                             Ok((Some(spec), spec_resp.compute_ctl_config))
     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) = &result {
     125            0 :             error!("attempt {} to get spec failed with: {}", attempt, e);
     126              :         } else {
     127            0 :             return result;
     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 :     Err(anyhow::anyhow!(
     136            0 :         "Exhausted all attempts to retrieve the spec from the control plane"
     137            0 :     ))
     138            0 : }
     139              : 
     140              : /// Check `pg_hba.conf` and update if needed to allow external connections.
     141            0 : pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
     142            0 :     // XXX: consider making it a part of spec.json
     143            0 :     let pghba_path = pgdata_path.join("pg_hba.conf");
     144            0 : 
     145            0 :     if config::line_in_file(&pghba_path, PG_HBA_ALL_MD5)? {
     146            0 :         info!("updated pg_hba.conf to allow external connections");
     147              :     } else {
     148            0 :         info!("pg_hba.conf is up-to-date");
     149              :     }
     150              : 
     151            0 :     Ok(())
     152            0 : }
     153              : 
     154              : /// Create a standby.signal file
     155            0 : pub fn add_standby_signal(pgdata_path: &Path) -> Result<()> {
     156            0 :     // XXX: consider making it a part of spec.json
     157            0 :     let signalfile = pgdata_path.join("standby.signal");
     158            0 : 
     159            0 :     if !signalfile.exists() {
     160            0 :         File::create(signalfile)?;
     161            0 :         info!("created standby.signal");
     162              :     } else {
     163            0 :         info!("reused pre-existing standby.signal");
     164              :     }
     165            0 :     Ok(())
     166            0 : }
     167              : 
     168              : #[instrument(skip_all)]
     169              : pub async fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
     170              :     let query = "ALTER EXTENSION neon UPDATE";
     171              :     info!("update neon extension version with query: {}", query);
     172              :     client.simple_query(query).await?;
     173              : 
     174              :     Ok(())
     175              : }
     176              : 
     177              : #[instrument(skip_all)]
     178              : pub async fn handle_migrations(client: &mut Client) -> Result<()> {
     179              :     info!("handle migrations");
     180              : 
     181              :     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
     182              :     // !BE SURE TO ONLY ADD MIGRATIONS TO THE END OF THIS ARRAY. IF YOU DO NOT, VERY VERY BAD THINGS MAY HAPPEN!
     183              :     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
     184              : 
     185              :     // Add new migrations in numerical order.
     186              :     let migrations = [
     187              :         include_str!("./migrations/0001-neon_superuser_bypass_rls.sql"),
     188              :         include_str!("./migrations/0002-alter_roles.sql"),
     189              :         include_str!("./migrations/0003-grant_pg_create_subscription_to_neon_superuser.sql"),
     190              :         include_str!("./migrations/0004-grant_pg_monitor_to_neon_superuser.sql"),
     191              :         include_str!("./migrations/0005-grant_all_on_tables_to_neon_superuser.sql"),
     192              :         include_str!("./migrations/0006-grant_all_on_sequences_to_neon_superuser.sql"),
     193              :         include_str!(
     194              :             "./migrations/0007-grant_all_on_tables_to_neon_superuser_with_grant_option.sql"
     195              :         ),
     196              :         include_str!(
     197              :             "./migrations/0008-grant_all_on_sequences_to_neon_superuser_with_grant_option.sql"
     198              :         ),
     199              :         include_str!("./migrations/0009-revoke_replication_for_previously_allowed_roles.sql"),
     200              :         include_str!(
     201              :             "./migrations/0010-grant_snapshot_synchronization_funcs_to_neon_superuser.sql"
     202              :         ),
     203              :         include_str!(
     204              :             "./migrations/0011-grant_pg_show_replication_origin_status_to_neon_superuser.sql"
     205              :         ),
     206              :     ];
     207              : 
     208              :     MigrationRunner::new(client, &migrations)
     209              :         .run_migrations()
     210              :         .await?;
     211              : 
     212              :     Ok(())
     213              : }
        

Generated by: LCOV version 2.1-beta