From 04f5807b998d7ffdeec1c8f0d6738f6676caa2af Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Thu, 10 Oct 2024 11:32:36 +0300 Subject: [PATCH] Add endpoint that allows extensions to be installed --- compute_tools/src/compute.rs | 23 ++++++++++++++++ compute_tools/src/http/api.rs | 26 ++++++++++++++++++ compute_tools/src/http/openapi_spec.yaml | 35 ++++++++++++++++++++++++ 3 files changed, 84 insertions(+) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 285be56264..76a0a566b4 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -19,6 +19,7 @@ use futures::future::join_all; use futures::stream::FuturesUnordered; use futures::StreamExt; use nix::unistd::Pid; +use postgres::config::Config; use postgres::error::SqlState; use postgres::{Client, NoTls}; use tracing::{debug, error, info, instrument, warn}; @@ -1367,6 +1368,28 @@ LIMIT 100", download_size } + pub fn install_extension(&self, ext_name: &str, db_name: &str) -> Result<()> { + let mut conf = + Config::from_str(self.connstr.as_str()).context("Failed to parse connection string")?; + conf.dbname(db_name); + + let mut db_client = conf + .connect(NoTls) + .context("Failed to connect to the database")?; + + let query = format!( + "CREATE EXTENSION IF NOT EXISTS {}", + ext_name.to_string().pg_quote() + ); + info!("creating extension with query: {}", query); + + db_client + .execute(&query, &[]) + .context(format!("Failed to execute query: {}", query))?; + + Ok(()) + } + #[tokio::main] pub async fn prepare_preload_libraries( &self, diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 79e6158081..b70c1ac802 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -98,6 +98,32 @@ async fn routes(req: Request, compute: &Arc) -> Response { + info!("serving /extensions POST request"); + let status = compute.get_status(); + if status != ComputeStatus::Running { + let msg = format!( + "invalid compute status for extensions request: {:?}", + status + ); + error!(msg); + return Response::new(Body::from(msg)); + } + + let body = hyper::body::to_bytes(req.into_body()).await.unwrap(); + let body = serde_json::from_slice::(&body).unwrap(); + let extension = body["extension"].as_str().unwrap(); + let database = body["database"].as_str().unwrap(); + let res = compute.install_extension(extension, database); + match res { + Ok(_) => Response::new(Body::from("true")), + Err(e) => { + error!("install_extension failed: {}", e); + Response::new(Body::from(e.to_string())) + } + } + } + (&Method::GET, "/info") => { let num_cpus = num_cpus::get_physical(); info!("serving /info GET request. num_cpus: {}", num_cpus); diff --git a/compute_tools/src/http/openapi_spec.yaml b/compute_tools/src/http/openapi_spec.yaml index e9fa66b323..600539d33f 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -144,6 +144,41 @@ paths: description: Error text or 'true' if check passed. example: "true" + /extensions: + post: + tags: + - Extensions + summary: Install extension if possible. + description: "" + operationId: installExtension + requestBody: + description: Extension name and database to install it to. + required: true + content: + application/json: + schema: + type: object + required: + - extension + - database + properties: + extension: + type: string + description: Extension name. + database: + type: string + description: Database name. + example: "neondb" + responses: + 200: + description: Result from extension installation + content: + text/plain: + schema: + type: string + description: Error text or 'true' if extension was installed. + example: "true" + /configure: post: tags: