Add extension neon to all databases.

- Run CREATE EXTENSION neon for template1, so that it was created in all databases.
- Run ALTER EXTENSION neon in all databases, to always have the newest version of the extension in computes.
- Add test_neon_extension test

compute_tools/src/compute.rs

@@ -693,11 +693,13 @@ 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())?;
         create_availability_check_data(&mut client)?;
 
         // 'Close' connection
@@ -737,11 +739,13 @@ 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())?;
         }
 
         // 'Close' connection

compute_tools/src/spec.rs

@@ -674,3 +674,72 @@ 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
+#[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 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)?;
+    }
+
+    Ok(())
+}

test_runner/regress/test_backpressure.py

@@ -24,8 +24,6 @@ def check_backpressure(endpoint: Endpoint, stop_event: threading.Event, polling_
     log.info("checks started")
 
     with pg_cur(endpoint) as cur:
-        cur.execute("CREATE EXTENSION neon")  # TODO move it to neon_fixtures?
-
         cur.execute("select pg_size_bytes(current_setting('max_replication_write_lag'))")
         res = cur.fetchone()
         max_replication_write_lag_bytes = res[0]
@@ -102,9 +100,13 @@ def test_backpressure_received_lsn_lag(neon_env_builder: NeonEnvBuilder):
     # Create a branch for us
     env.neon_cli.create_branch("test_backpressure")
 
-    endpoint = env.endpoints.create_start(
+    endpoint = env.endpoints.create(
         "test_backpressure", config_lines=["max_replication_write_lag=30MB"]
     )
+    # don't skip pg_catalog updates - it runs CREATE EXTENSION neon
+    # which is needed for backpressure_lsns() to work
+    endpoint.respec(skip_pg_catalog_updates=False)
+    endpoint.start()
     log.info("postgres is running on 'test_backpressure' branch")
 
     # setup check thread

test_runner/regress/test_neon_extension.py

@@ -0,0 +1,38 @@
+from contextlib import closing
+
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import NeonEnvBuilder
+
+
+# Verify that the neon extension is installed and has the correct version.
+def test_neon_extension(neon_env_builder: NeonEnvBuilder):
+    env = neon_env_builder.init_start()
+    env.neon_cli.create_branch("test_create_extension_neon")
+
+    endpoint_main = env.endpoints.create("test_create_extension_neon")
+    # don't skip pg_catalog updates - it runs CREATE EXTENSION neon
+    endpoint_main.respec(skip_pg_catalog_updates=False)
+    endpoint_main.start()
+
+    log.info("postgres is running on 'test_create_extension_neon' branch")
+
+    with closing(endpoint_main.connect()) as conn:
+        with conn.cursor() as cur:
+            cur.execute("SELECT extversion from pg_extension where extname='neon'")
+            # If this fails, it means the extension is either not installed
+            # or was updated and the version is different.
+            #
+            # IMPORTANT:
+            # 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.0",)
+
+            # 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.0",)

test_runner/regress/test_timeline_size.py

@@ -152,17 +152,20 @@ def test_timeline_size_quota(neon_env_builder: NeonEnvBuilder):
 
     wait_for_timeline_size_init(client, tenant=env.initial_tenant, timeline=new_timeline_id)
 
-    endpoint_main = env.endpoints.create_start(
+    endpoint_main = env.endpoints.create(
         "test_timeline_size_quota",
         # Set small limit for the test
         config_lines=["neon.max_cluster_size=30MB"],
     )
+    # don't skip pg_catalog updates - it runs CREATE EXTENSION neon
+    # which is needed for pg_cluster_size() to work
+    endpoint_main.respec(skip_pg_catalog_updates=False)
+    endpoint_main.start()
+
     log.info("postgres is running on 'test_timeline_size_quota' branch")
 
     with closing(endpoint_main.connect()) as conn:
         with conn.cursor() as cur:
-            cur.execute("CREATE EXTENSION neon")  # TODO move it to neon_fixtures?
-
             cur.execute("CREATE TABLE foo (t text)")
 
             wait_for_pageserver_catchup(endpoint_main)