Line data    Source code

       1              : //! FIXME: most of this is copy-paste from mgmt_api.rs ; dedupe into a `reqwest_utils::Client` crate.
       2              : use pageserver_client::mgmt_api::{Error, ResponseErrorMessageExt};
       3              : use reqwest::Method;
       4              : use serde::{Deserialize, Serialize};
       5              : use tokio_util::sync::CancellationToken;
       6              : use tracing::error;
       7              : 
       8              : use super::importbucket_format::Spec;
       9              : use crate::config::PageServerConf;
      10              : 
      11              : pub struct Client {
      12              :     base_url: String,
      13              :     authorization_header: Option<String>,
      14              :     client: reqwest::Client,
      15              :     cancel: CancellationToken,
      16              : }
      17              : 
      18              : pub type Result<T> = std::result::Result<T, Error>;
      19              : 
      20            0 : #[derive(Serialize, Deserialize, Debug)]
      21              : struct ImportProgressRequest {
      22              :     // no fields yet, not sure if there every will be any
      23              : }
      24              : 
      25            0 : #[derive(Serialize, Deserialize, Debug)]
      26              : struct ImportProgressResponse {
      27              :     // we don't care
      28              : }
      29              : 
      30              : impl Client {
      31            0 :     pub fn new(conf: &PageServerConf, cancel: CancellationToken) -> anyhow::Result<Self> {
      32            0 :         let Some(ref base_url) = conf.import_pgdata_upcall_api else {
      33            0 :             anyhow::bail!("import_pgdata_upcall_api is not configured")
      34              :         };
      35            0 :         Ok(Self {
      36            0 :             base_url: base_url.to_string(),
      37            0 :             client: reqwest::Client::new(),
      38            0 :             cancel,
      39            0 :             authorization_header: conf
      40            0 :                 .import_pgdata_upcall_api_token
      41            0 :                 .as_ref()
      42            0 :                 .map(|secret_string| secret_string.get_contents())
      43            0 :                 .map(|jwt| format!("Bearer {jwt}")),
      44            0 :         })
      45            0 :     }
      46              : 
      47            0 :     fn start_request<U: reqwest::IntoUrl>(
      48            0 :         &self,
      49            0 :         method: Method,
      50            0 :         uri: U,
      51            0 :     ) -> reqwest::RequestBuilder {
      52            0 :         let req = self.client.request(method, uri);
      53            0 :         if let Some(value) = &self.authorization_header {
      54            0 :             req.header(reqwest::header::AUTHORIZATION, value)
      55              :         } else {
      56            0 :             req
      57              :         }
      58            0 :     }
      59              : 
      60            0 :     async fn request_noerror<B: serde::Serialize, U: reqwest::IntoUrl>(
      61            0 :         &self,
      62            0 :         method: Method,
      63            0 :         uri: U,
      64            0 :         body: B,
      65            0 :     ) -> Result<reqwest::Response> {
      66            0 :         self.start_request(method, uri)
      67            0 :             .json(&body)
      68            0 :             .send()
      69            0 :             .await
      70            0 :             .map_err(Error::ReceiveBody)
      71            0 :     }
      72              : 
      73            0 :     async fn request<B: serde::Serialize, U: reqwest::IntoUrl>(
      74            0 :         &self,
      75            0 :         method: Method,
      76            0 :         uri: U,
      77            0 :         body: B,
      78            0 :     ) -> Result<reqwest::Response> {
      79            0 :         let res = self.request_noerror(method, uri, body).await?;
      80            0 :         let response = res.error_from_body().await?;
      81            0 :         Ok(response)
      82            0 :     }
      83              : 
      84            0 :     pub async fn send_progress_once(&self, spec: &Spec) -> Result<()> {
      85            0 :         let url = format!(
      86            0 :             "{}/projects/{}/branches/{}/import_progress",
      87            0 :             self.base_url, spec.project_id, spec.branch_id
      88            0 :         );
      89            0 :         let ImportProgressResponse {} = self
      90            0 :             .request(Method::POST, url, &ImportProgressRequest {})
      91            0 :             .await?
      92            0 :             .json()
      93            0 :             .await
      94            0 :             .map_err(Error::ReceiveBody)?;
      95            0 :         Ok(())
      96            0 :     }
      97              : 
      98            0 :     pub async fn send_progress_until_success(&self, spec: &Spec) -> anyhow::Result<()> {
      99              :         loop {
     100            0 :             match self.send_progress_once(spec).await {
     101            0 :                 Ok(()) => return Ok(()),
     102            0 :                 Err(Error::Cancelled) => return Err(anyhow::anyhow!("cancelled")),
     103            0 :                 Err(err) => {
     104            0 :                     error!(?err, "error sending progress, retrying");
     105            0 :                     if tokio::time::timeout(
     106            0 :                         std::time::Duration::from_secs(10),
     107            0 :                         self.cancel.cancelled(),
     108            0 :                     )
     109            0 :                     .await
     110            0 :                     .is_ok()
     111              :                     {
     112            0 :                         anyhow::bail!("cancelled while sending early progress update");
     113            0 :                     }
     114              :                 }
     115              :             }
     116              :         }
     117            0 :     }
     118              : }