From 87b8ac3ec38a3d275001d8471403567f0ee7a9cf Mon Sep 17 00:00:00 2001 From: Anastasia Lubennikova Date: Sun, 26 Nov 2023 08:37:01 +0000 Subject: [PATCH] Only create neon extension in postgres database; (#5918) Create neon extension in neon schema. --- compute_tools/src/compute.rs | 6 +- compute_tools/src/spec.rs | 85 +++++----------------- test_runner/regress/test_neon_extension.py | 10 --- test_runner/regress/test_timeline_size.py | 9 ++- 4 files changed, 25 insertions(+), 85 deletions(-) diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 0a77c3922a..5ace8ca1d2 100644 --- 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 index ae705702c5..97aa144c79 100644 --- 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)?; +pub fn handle_extension_neon(client: &mut Client) -> Result<()> { + info!("handle extension neon"); - // 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 query = "CREATE SCHEMA IF NOT EXISTS neon"; + client.simple_query(query)?; - let mut conf = Config::from_str(connstr)?; - conf.dbname(&db.name); + query = "CREATE EXTENSION IF NOT EXISTS neon WITH SCHEMA neon"; + info!("create neon extension with query: {}", query); + client.simple_query(query)?; - let mut db_client = conf.connect(NoTls)?; + 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 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)?; - } + // 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 index 1cd1c1a3d4..998f84f968 100644 --- 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 index 1146ca4b82..eb98348823 100644 --- 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}")