From ddbbac5a21c4b5652680c9e40b579faaa02b7847 Mon Sep 17 00:00:00 2001
From: Anastasia Lubennikova <anastasia@neon.tech>
Date: Fri, 5 Jan 2024 16:22:39 +0000
Subject: [PATCH] pre-install anon extension, if needed and GRANT all
 priveleged needed to use it to db_owner

---
 compute_tools/src/spec.rs | 86 +++++++++++++++++++++++++++++++++++----
 1 file changed, 78 insertions(+), 8 deletions(-)

diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs
index 681c89b8cdd0a..94f80ad5f6eee 100644
--- a/compute_tools/src/spec.rs
+++ b/compute_tools/src/spec.rs
@@ -541,6 +541,8 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
                 );
                 let _guard = info_span!("executing", query).entered();
                 client.execute(query.as_str(), &[])?;
+                // grants only
+                handle_extension_anon(spec, client, &db.owner.pg_quote(), true)?;
             }
             DatabaseAction::Create => {
                 let mut query: String = format!("CREATE DATABASE {} ", name.pg_quote());
@@ -552,6 +554,9 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
                     name.pg_quote()
                 );
                 client.execute(grant_query.as_str(), &[])?;
+
+                // Create anon extension if this compute needs it
+                handle_extension_anon(spec, client, &db.owner.pg_quote(), false)?;
             }
         };
 
@@ -679,14 +684,6 @@ pub fn handle_extensions(spec: &ComputeSpec, client: &mut Client) -> Result<()>
             info!("creating system extensions with query: {}", query);
             client.simple_query(query)?;
         }
-
-        if libs.contains("anon") {
-            // Create anon extension if this compute needs it
-            // Users cannot create it themselves, because superuser is required
-            let query = "CREATE EXTENSION IF NOT EXISTS anon";
-            info!("creating anon extension with query: {}", query);
-            client.simple_query(query)?;
-        }
     }
 
     Ok(())
@@ -721,3 +718,76 @@ pub fn handle_extension_neon(client: &mut Client) -> Result<()> {
 
     Ok(())
 }
+
+/// Connect to the database as superuser and pre-create anon extension
+/// if it is present in shared_preload_libraries
+#[instrument(skip_all)]
+pub fn handle_extension_anon(
+    spec: &ComputeSpec,
+    client: &mut Client,
+    db_owner: &str,
+    grants_only: bool,
+) -> Result<()> {
+    info!("handle extension anon");
+
+    if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
+        if libs.contains("anon") {
+            if !grants_only {
+                // check if extension is already initialized using anon.is_initialized()
+                let query = "SELECT anon.is_initialized()";
+                // TODO handle error
+                let rows = client.query(query, &[])?;
+                if rows.len() > 0 {
+                    let is_initialized: bool = rows[0].get(0);
+                    if is_initialized {
+                        info!("anon extension is already initialized");
+                        return Ok(());
+                    }
+                }
+
+                // Create anon extension if this compute needs it
+                // Users cannot create it themselves, because superuser is required.
+                let mut query = "CREATE EXTENSION IF NOT EXISTS anon CASCADE";
+                info!("creating anon extension with query: {}", query);
+                client.simple_query(query)?;
+
+                // Initialize anon extension
+                // This also requires superuser privileges, so users cannot do it themselves.
+                query = "SELECT anon.init()";
+                info!("initializing anon extension with query: {}", query);
+                client.simple_query(query)?;
+            }
+
+            // Grant permissions to db_owner to use anon extension functions
+            let query = format!("GRANT ALL ON ALL FUNCTIONS IN SCHEMA anon TO {}", db_owner);
+            info!("granting anon extension permissions with query: {}", query);
+            client.simple_query(&query)?;
+
+            // This is needed, because some functions are defined as SECURITY DEFINER.
+            // In Postgres SECURITY DEFINER functions are executed with the privileges
+            // of the owner.
+            // In anon extension this it is needed to access some GUCs, which are only accessible to
+            // superuser. But we've patched postgres to allow db_owner to access them as well.
+            // So we need to change owner of these functions to db_owner.
+            let query = format!("
+                SELECT 'ALTER FUNCTION '||nsp.nspname||'.'||p.proname||'('||pg_get_function_identity_arguments(p.oid)||') OWNER TO {};'
+                from pg_proc p
+                join pg_namespace nsp ON p.pronamespace = nsp.oid
+                where nsp.nspname = 'anon';", db_owner);
+
+            info!("change anon extension functions owner to db owner");
+            client.simple_query(&query)?;
+
+            //  affects views as well
+            let query = format!("GRANT ALL ON ALL TABLES IN SCHEMA anon TO {}", db_owner);
+            info!("granting anon extension permissions with query: {}", query);
+            client.simple_query(&query)?;
+
+            let query = format!("GRANT ALL ON ALL SEQUENCES IN SCHEMA anon TO {}", db_owner);
+            info!("granting anon extension permissions with query: {}", query);
+            client.simple_query(&query)?;
+        }
+    }
+
+    Ok(())
+}