Only create neon extension in postgres database; (#5918)

Create neon extension in neon schema.
neondatabase · Nov 26, 2023 · 87b8ac3 · 87b8ac3 · github-actions · Nov 26, 2023
1 parent 6b1c4cc
commit 87b8ac3
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 89 deletions.
diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs
@@ -693,13 +693,12 @@ impl ComputeNode {
         let spec = &compute_state.pspec.as_ref().expect("spec must be set").spec;
         create_neon_superuser(spec, &mut client)?;
         cleanup_instance(&mut client)?;
-        handle_extension_neon(self.connstr.as_str())?;
         handle_roles(spec, &mut client)?;
         handle_databases(spec, &mut client)?;
         handle_role_deletions(spec, self.connstr.as_str(), &mut client)?;
         handle_grants(spec, &mut client, self.connstr.as_str())?;
         handle_extensions(spec, &mut client)?;
-        handle_alter_extension_neon(spec, &mut client, self.connstr.as_str())?;
+        handle_extension_neon(&mut client)?;
         create_availability_check_data(&mut client)?;
 
         // 'Close' connection
@@ -739,13 +738,12 @@ impl ComputeNode {
         if spec.mode == ComputeMode::Primary {
             client.simple_query("SET neon.forward_ddl = false")?;
             cleanup_instance(&mut client)?;
-            handle_extension_neon(self.connstr.as_str())?;
             handle_roles(&spec, &mut client)?;
             handle_databases(&spec, &mut client)?;
             handle_role_deletions(&spec, self.connstr.as_str(), &mut client)?;
             handle_grants(&spec, &mut client, self.connstr.as_str())?;
             handle_extensions(&spec, &mut client)?;
-            handle_alter_extension_neon(&spec, &mut client, self.connstr.as_str())?;
+            handle_extension_neon(&mut client)?;
         }
 
         // 'Close' connection

diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs
@@ -675,78 +675,29 @@ pub fn handle_extensions(spec: &ComputeSpec, client: &mut Client) -> Result<()>
     Ok(())
 }
 
-/// connect to template1 and postgres to create neon extension
-/// that will be available in all databases
-pub fn handle_extension_neon(connstr: &str) -> Result<()> {
-    for dbname in ["template1", "postgres"].iter() {
-        let mut conf = Config::from_str(connstr)?;
-        conf.dbname(dbname);
-        let mut template1_client = conf.connect(NoTls)?;
-
-        let create_extension_neon_query = "CREATE EXTENSION IF NOT EXISTS neon";
-        info!(
-            "creating neon extension with query: {} in db {}",
-            create_extension_neon_query, dbname
-        );
-        template1_client.simple_query(create_extension_neon_query)?;
-    }
-
-    Ok(())
-}
-
-/// Run ALTER EXTENSION neon UPDATE for each valid database
+/// Run CREATE and ALTER EXTENSION neon UPDATE for postgres database
 #[instrument(skip_all)]
-pub fn handle_alter_extension_neon(
-    spec: &ComputeSpec,
-    client: &mut Client,
-    connstr: &str,
-) -> Result<()> {
-    info!("modifying database permissions");
-    let existing_dbs = get_existing_dbs(client)?;
-
-    // We'd better do this at db creation time, but we don't always know when it happens.
-    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 grants 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)?;
-
-        // this will be a no-op if extension is already up to date,
-        // which may happen in two cases:
-        // - extension was just installed
-        // - extension was already installed and is up to date
-        let create_extension_neon_query = "CREATE EXTENSION IF NOT EXISTS neon";
-        info!(
-            "create extension neon query for db {} : {}",
-            &db.name, &create_extension_neon_query
-        );
-        db_client.simple_query(create_extension_neon_query)?;
-
-        let alter_extension_neon_query = "ALTER EXTENSION neon UPDATE";
-        info!(
-            "alter extension neon query for db {} : {}",
-            &db.name, &alter_extension_neon_query
-        );
-        db_client.simple_query(alter_extension_neon_query)?;
-    }
+pub fn handle_extension_neon(client: &mut Client) -> Result<()> {
+    info!("handle extension neon");
+
+    let mut query = "CREATE SCHEMA IF NOT EXISTS neon";
+    client.simple_query(query)?;
+
+    query = "CREATE EXTENSION IF NOT EXISTS neon WITH SCHEMA neon";
+    info!("create neon extension with query: {}", query);
+    client.simple_query(query)?;
+
+    query = "ALTER EXTENSION neon SET SCHEMA neon";
+    info!("alter neon extension schema with query: {}", query);
+    client.simple_query(query)?;
+
+    // this will be a no-op if extension is already up to date,
+    // which may happen in two cases:
+    // - extension was just installed
+    // - extension was already installed and is up to date
+    let query = "ALTER EXTENSION neon UPDATE";
+    info!("update neon extension schema with query: {}", query);
+    client.simple_query(query)?;
 
     Ok(())
 }
diff --git a/test_runner/regress/test_neon_extension.py b/test_runner/regress/test_neon_extension.py
@@ -26,13 +26,3 @@ def test_neon_extension(neon_env_builder: NeonEnvBuilder):
             # If the version has changed, the test should be updated.
             # Ensure that the default version is also updated in the neon.control file
             assert cur.fetchone() == ("1.1",)
-
-            # create test database
-            cur.execute("CREATE DATABASE foodb")
-
-    # connect to test database
-    # test that neon extension is installed and has the correct version
-    with closing(endpoint_main.connect(dbname="foodb")) as conn:
-        with conn.cursor() as cur:
-            cur.execute("SELECT extversion from pg_extension where extname='neon'")
-            assert cur.fetchone() == ("1.1",)
diff --git a/test_runner/regress/test_timeline_size.py b/test_runner/regress/test_timeline_size.py
@@ -134,10 +134,11 @@ def wait_for_pageserver_catchup(endpoint_main: Endpoint, polling_interval=1, tim
         res = endpoint_main.safe_psql(
             """
             SELECT
-                pg_size_pretty(pg_cluster_size()),
+                pg_size_pretty(neon.pg_cluster_size()),
                 pg_wal_lsn_diff(pg_current_wal_flush_lsn(), received_lsn) as received_lsn_lag
-            FROM backpressure_lsns();
-            """
+            FROM neon.backpressure_lsns();
+            """,
+            dbname="postgres",
         )[0]
         log.info(f"pg_cluster_size = {res[0]}, received_lsn_lag = {res[1]}")
         received_lsn_lag = res[1]
@@ -214,7 +215,7 @@ def test_timeline_size_quota(neon_env_builder: NeonEnvBuilder):
 
             wait_for_pageserver_catchup(endpoint_main)
 
-            cur.execute("SELECT * from pg_size_pretty(pg_cluster_size())")
+            cur.execute("SELECT * from pg_size_pretty(neon.pg_cluster_size())")
             pg_cluster_size = cur.fetchone()
             log.info(f"pg_cluster_size = {pg_cluster_size}")