From 028f5291f813006eac710095733c281adc1aa59d Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Wed, 16 Oct 2024 11:55:37 +0100
Subject: [PATCH] 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<String> = db_client
-            .query_one(version_query, &[&ext_name])
+        let version: Option<String> = 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