From 04f5807b998d7ffdeec1c8f0d6738f6676caa2af Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Thu, 10 Oct 2024 11:32:36 +0300 Subject: [PATCH 01/10] 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: From a766ac3dc7d90e04c58ab15a0d653926f841d642 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Thu, 10 Oct 2024 12:37:40 +0300 Subject: [PATCH 02/10] Add error response status code --- compute_tools/src/http/api.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index b70c1ac802..83e95652b2 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -119,7 +119,7 @@ async fn routes(req: Request, compute: &Arc) -> Response Response::new(Body::from("true")), Err(e) => { error!("install_extension failed: {}", e); - Response::new(Body::from(e.to_string())) + render_json_error(&e.to_string(), StatusCode::INTERNAL_SERVER_ERROR) } } } From 2b5d79f625604ac24f85ee6e1253ed353f6a9930 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Thu, 10 Oct 2024 12:51:53 +0300 Subject: [PATCH 03/10] Respond with extension name and version --- compute_tools/src/compute.rs | 14 ++++++++++++-- compute_tools/src/http/api.rs | 9 ++++++++- compute_tools/src/http/openapi_spec.yaml | 22 ++++++++++++++++++---- libs/compute_api/src/responses.rs | 6 ++++++ 4 files changed, 44 insertions(+), 7 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 76a0a566b4..036dd39a49 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1368,7 +1368,7 @@ LIMIT 100", download_size } - pub fn install_extension(&self, ext_name: &str, db_name: &str) -> Result<()> { + 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); @@ -1387,7 +1387,17 @@ LIMIT 100", .execute(&query, &[]) .context(format!("Failed to execute query: {}", query))?; - Ok(()) + let version_query = format!( + "SELECT extversion FROM pg_extension WHERE extname = '{}'", + ext_name.to_string().pg_quote() + ); + + let version: String = db_client + .query_one(&version_query, &[]) + .context(format!("Failed to execute query: {}", version_query))? + .get(0); + + Ok(version) } #[tokio::main] diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 83e95652b2..947b91a311 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -10,6 +10,7 @@ use crate::catalog::{get_database_schema, get_dbs_and_roles}; use crate::compute::forward_termination_signal; use crate::compute::{ComputeNode, ComputeState, ParsedSpec}; use compute_api::requests::ConfigurationRequest; +use compute_api::responses::ExtensionInstallResult; use compute_api::responses::{ComputeStatus, ComputeStatusResponse, GenericAPIError}; use anyhow::Result; @@ -116,7 +117,13 @@ async fn routes(req: Request, compute: &Arc) -> Response Response::new(Body::from("true")), + Ok(res) => render_json(Body::from( + serde_json::to_string(&ExtensionInstallResult { + extension: extension.to_string(), + version: res, + }) + .unwrap(), + )), Err(e) => { error!("install_extension failed: {}", e); render_json_error(&e.to_string(), StatusCode::INTERNAL_SERVER_ERROR) diff --git a/compute_tools/src/http/openapi_spec.yaml b/compute_tools/src/http/openapi_spec.yaml index 600539d33f..9af406d38f 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -173,11 +173,15 @@ paths: 200: description: Result from extension installation content: - text/plain: + application/json: schema: - type: string - description: Error text or 'true' if extension was installed. - example: "true" + $ref: "#/components/schemas/ExtensionInstallResult" + 500: + description: Error during extension installation. + content: + application/json: + schema: + $ref: "#/components/schemas/GenericError" /configure: post: @@ -444,6 +448,16 @@ components: - configuration example: running + ExtensionInstallResult: + type: object + properties: + extension: + description: Name of the extension. + type: string + version: + description: Version of the extension. + type: string + InstalledExtensions: type: object properties: diff --git a/libs/compute_api/src/responses.rs b/libs/compute_api/src/responses.rs index 5023fce003..7d96975f3f 100644 --- a/libs/compute_api/src/responses.rs +++ b/libs/compute_api/src/responses.rs @@ -168,3 +168,9 @@ pub struct InstalledExtension { pub struct InstalledExtensions { pub extensions: Vec, } + +#[derive(Clone, Debug, Default, Serialize)] +pub struct ExtensionInstallResult { + pub extension: String, + pub version: String, +} From 43400dfdeaa53722db17cbd8363317a0d7ae157f Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Thu, 10 Oct 2024 13:34:56 +0300 Subject: [PATCH 04/10] change path --- compute_tools/src/http/api.rs | 4 ++-- compute_tools/src/http/openapi_spec.yaml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 947b91a311..7542f5ef36 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -99,8 +99,8 @@ async fn routes(req: Request, compute: &Arc) -> Response { - info!("serving /extensions POST request"); + (&Method::POST, "/extensions/install") => { + info!("serving /extensions/install POST request"); let status = compute.get_status(); if status != ComputeStatus::Running { let msg = format!( diff --git a/compute_tools/src/http/openapi_spec.yaml b/compute_tools/src/http/openapi_spec.yaml index 9af406d38f..9bf5aeebf1 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -144,7 +144,7 @@ paths: description: Error text or 'true' if check passed. example: "true" - /extensions: + /extensions/install: post: tags: - Extensions From 133ea2f31fc0b000245cf9865d5c6268603b79e1 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Fri, 11 Oct 2024 09:19:24 +0300 Subject: [PATCH 05/10] Require extension version parameter --- compute_tools/src/compute.rs | 27 +++++------ compute_tools/src/http/api.rs | 18 ++++---- compute_tools/src/http/openapi_spec.yaml | 57 ++++++++++++++---------- libs/compute_api/src/requests.rs | 7 +++ 4 files changed, 61 insertions(+), 48 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 036dd39a49..95ddf20901 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1368,32 +1368,27 @@ 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")?; + pub fn install_extension( + &self, + ext_name: &str, + db_name: &str, + ext_version: &str, + ) -> Result { + let mut conf = Config::from_str(self.connstr.as_str()).unwrap(); 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); - + let query = "CREATE EXTENSION IF NOT EXISTS $1 WITH VERSION $2"; db_client - .execute(&query, &[]) + .query(query, &[&ext_name, &ext_version]) .context(format!("Failed to execute query: {}", query))?; - let version_query = format!( - "SELECT extversion FROM pg_extension WHERE extname = '{}'", - ext_name.to_string().pg_quote() - ); - + let version_query = "SELECT extversion FROM pg_extension WHERE extname = $1"; let version: String = db_client - .query_one(&version_query, &[]) + .query_one(version_query, &[&ext_name]) .context(format!("Failed to execute query: {}", version_query))? .get(0); diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 7542f5ef36..f92e260071 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -9,9 +9,10 @@ use crate::catalog::SchemaDumpError; use crate::catalog::{get_database_schema, get_dbs_and_roles}; use crate::compute::forward_termination_signal; use crate::compute::{ComputeNode, ComputeState, ParsedSpec}; -use compute_api::requests::ConfigurationRequest; -use compute_api::responses::ExtensionInstallResult; -use compute_api::responses::{ComputeStatus, ComputeStatusResponse, GenericAPIError}; +use compute_api::requests::{ConfigurationRequest, ExtensionInstallRequest}; +use compute_api::responses::{ + ComputeStatus, ComputeStatusResponse, ExtensionInstallResult, GenericAPIError, +}; use anyhow::Result; use hyper::header::CONTENT_TYPE; @@ -111,15 +112,14 @@ async fn routes(req: Request, compute: &Arc) -> Response(&body).unwrap(); - let extension = body["extension"].as_str().unwrap(); - let database = body["database"].as_str().unwrap(); - let res = compute.install_extension(extension, database); + let request = hyper::body::to_bytes(req.into_body()).await.unwrap(); + let request = serde_json::from_slice::(&request).unwrap(); + let res = + compute.install_extension(&request.extension, &request.database, &request.version); match res { Ok(res) => render_json(Body::from( serde_json::to_string(&ExtensionInstallResult { - extension: extension.to_string(), + extension: request.extension, version: res, }) .unwrap(), diff --git a/compute_tools/src/http/openapi_spec.yaml b/compute_tools/src/http/openapi_spec.yaml index 9bf5aeebf1..90401ef2e5 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -10,7 +10,7 @@ paths: /status: get: tags: - - Info + - Info summary: Get compute node internal status. description: "" operationId: getComputeStatus @@ -25,7 +25,7 @@ paths: /metrics.json: get: tags: - - Info + - Info summary: Get compute node startup metrics in JSON format. description: "" operationId: getComputeMetricsJSON @@ -40,7 +40,7 @@ paths: /insights: get: tags: - - Info + - Info summary: Get current compute insights in JSON format. description: | Note, that this doesn't include any historical data. @@ -56,7 +56,7 @@ paths: /installed_extensions: get: tags: - - Info + - Info summary: Get installed extensions. description: "" operationId: getInstalledExtensions @@ -70,7 +70,7 @@ paths: /info: get: tags: - - Info + - Info summary: Get info about the compute pod / VM. description: "" operationId: getInfo @@ -130,7 +130,7 @@ paths: /check_writability: post: tags: - - Check + - Check summary: Check that we can write new data on this compute. description: "" operationId: checkComputeWritability @@ -147,7 +147,7 @@ paths: /extensions/install: post: tags: - - Extensions + - Extensions summary: Install extension if possible. description: "" operationId: installExtension @@ -157,18 +157,7 @@ paths: content: application/json: schema: - type: object - required: - - extension - - database - properties: - extension: - type: string - description: Extension name. - database: - type: string - description: Database name. - example: "neondb" + $ref: "#/components/schemas/ExtensionInstallRequest" responses: 200: description: Result from extension installation @@ -186,7 +175,7 @@ paths: /configure: post: tags: - - Configure + - Configure summary: Perform compute node configuration. description: | This is a blocking API endpoint, i.e. it blocks waiting until @@ -240,7 +229,7 @@ paths: /extension_server: post: tags: - - Extension + - Extension summary: Download extension from S3 to local folder. description: "" operationId: downloadExtension @@ -269,7 +258,7 @@ paths: /terminate: post: tags: - - Terminate + - Terminate summary: Terminate Postgres and wait for it to exit description: "" operationId: terminate @@ -408,7 +397,7 @@ components: moment, when spec was received. example: "2022-10-12T07:20:50.52Z" status: - $ref: '#/components/schemas/ComputeStatus' + $ref: "#/components/schemas/ComputeStatus" last_active: type: string description: | @@ -448,15 +437,37 @@ components: - configuration example: running + ExtensionInstallRequest: + type: object + required: + - extension + - database + - version + properties: + extension: + type: string + description: Extension name. + example: "pg_session_jwt" + version: + type: string + description: Version of the extension. + example: "1.0.0" + database: + type: string + description: Database name. + example: "neondb" + ExtensionInstallResult: type: object properties: extension: description: Name of the extension. type: string + example: "pg_session_jwt" version: description: Version of the extension. type: string + example: "1.0.0" InstalledExtensions: type: object diff --git a/libs/compute_api/src/requests.rs b/libs/compute_api/src/requests.rs index 5896c7dc65..34771834f3 100644 --- a/libs/compute_api/src/requests.rs +++ b/libs/compute_api/src/requests.rs @@ -12,3 +12,10 @@ use serde::Deserialize; pub struct ConfigurationRequest { pub spec: ComputeSpec, } + +#[derive(Deserialize, Debug)] +pub struct ExtensionInstallRequest { + pub extension: String, + pub database: String, + pub version: String, +} From 36bffe5ff39a65a7edf6fc4394a4c56f7fc967b5 Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Mon, 14 Oct 2024 18:35:17 +0300 Subject: [PATCH 06/10] Install required version or update --- compute_tools/src/compute.rs | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 95ddf20901..e8ead81d63 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1381,18 +1381,28 @@ LIMIT 100", .connect(NoTls) .context("Failed to connect to the database")?; - let query = "CREATE EXTENSION IF NOT EXISTS $1 WITH VERSION $2"; - db_client - .query(query, &[&ext_name, &ext_version]) - .context(format!("Failed to execute query: {}", query))?; - let version_query = "SELECT extversion FROM pg_extension WHERE extname = $1"; - let version: String = db_client + let version: Result = db_client .query_one(version_query, &[&ext_name]) - .context(format!("Failed to execute query: {}", version_query))? - .get(0); + .context(format!("Failed to execute query: {}", version_query)) + .map(|row| row.get(0)); - Ok(version) + if let Ok(installed_version) = version { + if installed_version == ext_version { + return Ok(installed_version); + } + let query = "ALTER EXTENSION $1 UPDATE TO $2"; + db_client + .execute(query, &[&ext_name, &ext_version]) + .context(format!("Failed to execute query: {}", query))?; + } else { + let query = "CREATE EXTENSION IF NOT EXISTS $1 WITH VERSION $2"; + db_client + .execute(query, &[&ext_name, &ext_version]) + .context(format!("Failed to execute query: {}", query))?; + } + + Ok(ext_version.to_string()) } #[tokio::main] From 603ca192a7db6283f454814400e169b9216d7aeb Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Mon, 14 Oct 2024 18:54:25 +0300 Subject: [PATCH 07/10] Change path as per Alexey's feedback --- compute_tools/src/http/api.rs | 4 ++-- compute_tools/src/http/openapi_spec.yaml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index f92e260071..2aae7ca400 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -100,8 +100,8 @@ async fn routes(req: Request, compute: &Arc) -> Response { - info!("serving /extensions/install POST request"); + (&Method::POST, "/extensions") => { + info!("serving /extensions POST request"); let status = compute.get_status(); if status != ComputeStatus::Running { let msg = format!( diff --git a/compute_tools/src/http/openapi_spec.yaml b/compute_tools/src/http/openapi_spec.yaml index 90401ef2e5..47952deda9 100644 --- a/compute_tools/src/http/openapi_spec.yaml +++ b/compute_tools/src/http/openapi_spec.yaml @@ -144,7 +144,7 @@ paths: description: Error text or 'true' if check passed. example: "true" - /extensions/install: + /extensions: post: tags: - Extensions From f105dc9397c9c89d3d9bf952fc15ad1a115e188d Mon Sep 17 00:00:00 2001 From: Jere Vaara Date: Wed, 16 Oct 2024 13:23:05 +0300 Subject: [PATCH 08/10] Make install extensions fn async --- compute_tools/src/compute.rs | 21 +++++++++++++++------ compute_tools/src/http/api.rs | 9 +++++---- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index e8ead81d63..fc442a77bf 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -1368,22 +1368,28 @@ LIMIT 100", download_size } - pub fn install_extension( + pub async fn install_extension( &self, ext_name: &str, db_name: &str, ext_version: &str, ) -> Result { + use tokio_postgres::config::Config; + use tokio_postgres::NoTls; + let mut conf = Config::from_str(self.connstr.as_str()).unwrap(); conf.dbname(db_name); - let mut db_client = conf + let (db_client, conn) = conf .connect(NoTls) + .await .context("Failed to connect to the database")?; + tokio::spawn(conn); let version_query = "SELECT extversion FROM pg_extension WHERE extname = $1"; let version: Result = db_client .query_one(version_query, &[&ext_name]) + .await .context(format!("Failed to execute query: {}", version_query)) .map(|row| row.get(0)); @@ -1391,14 +1397,17 @@ LIMIT 100", if installed_version == ext_version { return Ok(installed_version); } - let query = "ALTER EXTENSION $1 UPDATE TO $2"; + let query = format!("ALTER EXTENSION {ext_name:?} UPDATE TO {ext_version:?}"); db_client - .execute(query, &[&ext_name, &ext_version]) + .simple_query(&query) + .await .context(format!("Failed to execute query: {}", query))?; } else { - let query = "CREATE EXTENSION IF NOT EXISTS $1 WITH VERSION $2"; + let query = + format!("CREATE EXTENSION IF NOT EXISTS {ext_name:?} WITH VERSION {ext_version:?}"); db_client - .execute(query, &[&ext_name, &ext_version]) + .simple_query(&query) + .await .context(format!("Failed to execute query: {}", query))?; } diff --git a/compute_tools/src/http/api.rs b/compute_tools/src/http/api.rs index 2aae7ca400..4c43d0b598 100644 --- a/compute_tools/src/http/api.rs +++ b/compute_tools/src/http/api.rs @@ -114,13 +114,14 @@ async fn routes(req: Request, compute: &Arc) -> Response(&request).unwrap(); - let res = - compute.install_extension(&request.extension, &request.database, &request.version); + let res = compute + .install_extension(&request.extension, &request.database, &request.version) + .await; match res { - Ok(res) => render_json(Body::from( + Ok(version) => render_json(Body::from( serde_json::to_string(&ExtensionInstallResult { extension: request.extension, - version: res, + version, }) .unwrap(), )), From 028f5291f813006eac710095733c281adc1aa59d Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Wed, 16 Oct 2024 11:55:37 +0100 Subject: [PATCH 09/10] call out to quote_ident --- compute_tools/src/compute.rs | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index fc442a77bf..0d52d61585 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -19,8 +19,8 @@ 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::Row; use postgres::{Client, NoTls}; use tracing::{debug, error, info, instrument, warn}; use utils::id::{TenantId, TimelineId}; @@ -1387,24 +1387,34 @@ LIMIT 100", tokio::spawn(conn); let version_query = "SELECT extversion FROM pg_extension WHERE extname = $1"; - let version: Result = db_client - .query_one(version_query, &[&ext_name]) + let version: Option = db_client + .query_opt(version_query, &[&ext_name]) .await - .context(format!("Failed to execute query: {}", version_query)) + .with_context(|| format!("Failed to execute query: {}", version_query))? .map(|row| row.get(0)); - if let Ok(installed_version) = version { + // sanitize the inputs as postgres idents. + // should we instead use `Escaping::pg_quote`? + let format_query = "SELECT quote_ident($1) ext_name, quote_ident($2) ext_version"; + let format: Row = db_client + .query_one(format_query, &[&ext_name, &ext_version]) + .await + .with_context(|| format!("Failed to execute query: {}", format_query))?; + let ext_name: String = format.get("ext_name"); + let ext_version: String = format.get("ext_version"); + + if let Some(installed_version) = version { if installed_version == ext_version { return Ok(installed_version); } - let query = format!("ALTER EXTENSION {ext_name:?} UPDATE TO {ext_version:?}"); + let query = format!("ALTER EXTENSION {ext_name} UPDATE TO {ext_version}"); db_client .simple_query(&query) .await .context(format!("Failed to execute query: {}", query))?; } else { let query = - format!("CREATE EXTENSION IF NOT EXISTS {ext_name:?} WITH VERSION {ext_version:?}"); + format!("CREATE EXTENSION IF NOT EXISTS {ext_name} WITH VERSION {ext_version}"); db_client .simple_query(&query) .await From 59d197f6f6a3fbbc887f63e24431b775ee742892 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Wed, 16 Oct 2024 12:05:27 +0100 Subject: [PATCH 10/10] use pg_quote --- compute_tools/src/compute.rs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 0d52d61585..81d631fc0a 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -20,7 +20,6 @@ use futures::stream::FuturesUnordered; use futures::StreamExt; use nix::unistd::Pid; use postgres::error::SqlState; -use postgres::Row; use postgres::{Client, NoTls}; use tracing::{debug, error, info, instrument, warn}; use utils::id::{TenantId, TimelineId}; @@ -1394,14 +1393,8 @@ LIMIT 100", .map(|row| row.get(0)); // sanitize the inputs as postgres idents. - // should we instead use `Escaping::pg_quote`? - let format_query = "SELECT quote_ident($1) ext_name, quote_ident($2) ext_version"; - let format: Row = db_client - .query_one(format_query, &[&ext_name, &ext_version]) - .await - .with_context(|| format!("Failed to execute query: {}", format_query))?; - let ext_name: String = format.get("ext_name"); - let ext_version: String = format.get("ext_version"); + let ext_name: String = ext_name.to_string().pg_quote(); + let ext_version: String = ext_version.to_string().pg_quote(); if let Some(installed_version) = version { if installed_version == ext_version {