From ab72f94599012da638ff3f81e5fed5a07c6f641a Mon Sep 17 00:00:00 2001
From: Jere Vaara <jerevaara@gmail.com>
Date: Mon, 7 Oct 2024 16:04:20 +0300
Subject: [PATCH] create pg_session_jwt extension if local proxy spec has jwks

---
 compute_tools/src/compute.rs |  3 +++
 compute_tools/src/spec.rs    | 41 ++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
index 3e558a7d3c..e43f95fff0 100644
--- a/compute_tools/src/compute.rs
+++ b/compute_tools/src/compute.rs
@@ -890,6 +890,8 @@ impl ComputeNode {
         .context("apply_config handle_grants")?;
         handle_extensions(spec, &mut client).context("apply_config handle_extensions")?;
         handle_extension_neon(&mut client).context("apply_config handle_extension_neon")?;
+        handle_jwt_extension(spec, &mut client, connstr.as_str())
+            .context("apply_config handle_jwt_extension")?;
         create_availability_check_data(&mut client)
             .context("apply_config create_availability_check_data")?;
 
@@ -992,6 +994,7 @@ impl ComputeNode {
                 )?;
                 handle_extensions(&spec, &mut client)?;
                 handle_extension_neon(&mut client)?;
+                handle_jwt_extension(&spec, &mut client, self.connstr.as_str())?;
                 // We can skip handle_migrations here because a new migration can only appear
                 // if we have a new version of the compute_ctl binary, which can only happen
                 // if compute got restarted, in which case we'll end up inside of apply_config
diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs
index 73f3d1006a..fc7ae049a9 100644
--- a/compute_tools/src/spec.rs
+++ b/compute_tools/src/spec.rs
@@ -731,7 +731,48 @@ pub fn handle_extensions(spec: &ComputeSpec, client: &mut Client) -> Result<()>
             client.simple_query(query)?;
         }
     }
+    Ok(())
+}
 
+/// Create pg_session_jwt in all databases if configured
+#[instrument(skip_all)]
+pub fn handle_jwt_extension(spec: &ComputeSpec, client: &mut Client, connstr: &str) -> Result<()> {
+    if let Some(local_proxy) = &spec.local_proxy_config {
+        if let Some(jwks_list) = &local_proxy.jwks {
+            if !jwks_list.is_empty() {
+                info!("enabling pg_session_jwt extension");
+                let existing_dbs = get_existing_dbs(client)?;
+
+                for db in &spec.cluster.databases {
+                    match existing_dbs.get(&db.name) {
+                        Some(pg_db) => {
+                            if pg_db.restrict_conn || pg_db.invalid {
+                                info!(
+                                    "skipping extension for db {} (invalid: {}, connections not allowed: {})",
+                                    db.name, pg_db.invalid, pg_db.restrict_conn
+                                );
+                                continue;
+                            }
+                        }
+                        None => {
+                            bail!(
+                                "database {} doesn't exist in Postgres after handle_databases()",
+                                db.name
+                            );
+                        }
+                    }
+                    let mut conf = Config::from_str(connstr)?;
+                    conf.dbname(&db.name);
+
+                    let mut db_client = conf.connect(NoTls)?;
+
+                    let query = "CREATE EXTENSION IF NOT EXISTS pg_session_jwt";
+                    info!("creating pg_session_jwt extension with query: {}", query);
+                    db_client.simple_query(query)?;
+                }
+            }
+        }
+    }
     Ok(())
 }