From f8d9bd8d140d71e58ce8e4746a139df129635d7c Mon Sep 17 00:00:00 2001 From: Anastasia Lubennikova Date: Tue, 31 Oct 2023 13:28:59 +0000 Subject: [PATCH] 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 | 4 ++ compute_tools/src/spec.rs | 69 ++++++++++++++++++++++ test_runner/regress/test_backpressure.py | 8 ++- test_runner/regress/test_neon_extension.py | 38 ++++++++++++ test_runner/regress/test_timeline_size.py | 9 ++- 5 files changed, 122 insertions(+), 6 deletions(-) create mode 100644 test_runner/regress/test_neon_extension.py diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index 373f05ab2f..0a77c3922a 100644 --- a/compute_tools/src/compute.rs +++ b/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 diff --git a/compute_tools/src/spec.rs b/compute_tools/src/spec.rs index a85d6287b1..8a94950829 100644 --- a/compute_tools/src/spec.rs +++ b/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(()) +} diff --git a/test_runner/regress/test_backpressure.py b/test_runner/regress/test_backpressure.py index b14974279e..bc3faf9271 100644 --- a/test_runner/regress/test_backpressure.py +++ b/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 diff --git a/test_runner/regress/test_neon_extension.py b/test_runner/regress/test_neon_extension.py new file mode 100644 index 0000000000..21d01464bf --- /dev/null +++ b/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",) diff --git a/test_runner/regress/test_timeline_size.py b/test_runner/regress/test_timeline_size.py index c2e93c48c7..1146ca4b82 100644 --- a/test_runner/regress/test_timeline_size.py +++ b/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)