From 31aa0283b00078bcedc1f1d17a4a801bf4e752c6 Mon Sep 17 00:00:00 2001 From: Alek Westover Date: Fri, 23 Jun 2023 09:30:49 -0400 Subject: [PATCH] More Extension Features (#4555) Added tenant specific extensions and more tests --- compute_tools/src/bin/compute_ctl.rs | 23 +++- compute_tools/src/extension_server.rs | 5 +- test_runner/fixtures/neon_fixtures.py | 15 ++- test_runner/fixtures/types.py | 3 + .../regress/test_download_extensions.py | 115 ++++++++++++------ 5 files changed, 115 insertions(+), 46 deletions(-) diff --git a/compute_tools/src/bin/compute_ctl.rs b/compute_tools/src/bin/compute_ctl.rs index 708aab4a56..fae9ac9319 100644 --- a/compute_tools/src/bin/compute_ctl.rs +++ b/compute_tools/src/bin/compute_ctl.rs @@ -77,11 +77,10 @@ fn main() -> Result<()> { }; let rt = Runtime::new().unwrap(); - let copy_remote_storage = ext_remote_storage.clone(); - rt.block_on(async move { - download_extension(©_remote_storage, ExtensionType::Shared, pgbin) + rt.block_on(async { + download_extension(&ext_remote_storage, ExtensionType::Shared, pgbin) .await - .expect("download extension should work"); + .expect("download shared extensions should work"); }); let http_port = *matches @@ -182,14 +181,18 @@ fn main() -> Result<()> { } }; + dbg!(&spec); let mut new_state = ComputeState::new(); let spec_set; + let tenant_id; if let Some(spec) = spec { let pspec = ParsedSpec::try_from(spec).map_err(|msg| anyhow::anyhow!(msg))?; + tenant_id = Some(pspec.tenant_id.to_string()); new_state.pspec = Some(pspec); spec_set = true; } else { spec_set = false; + tenant_id = None; } let compute_node = ComputeNode { connstr: Url::parse(connstr).context("cannot parse connstr as a URL")?, @@ -198,7 +201,7 @@ fn main() -> Result<()> { live_config_allowed, state: Mutex::new(new_state), state_changed: Condvar::new(), - ext_remote_storage, + ext_remote_storage: ext_remote_storage.clone(), }; let compute = Arc::new(compute_node); @@ -224,6 +227,16 @@ fn main() -> Result<()> { } } + // Now we have the spec, so we request the tenant specific extensions + if let Some(tenant_id) = tenant_id { + let rt = Runtime::new().unwrap(); + rt.block_on(async { + download_extension(&ext_remote_storage, ExtensionType::Tenant(tenant_id), pgbin) + .await + .expect("download tenant specific extensions should work"); + }); + } + // We got all we need, update the state. let mut state = compute.state.lock().unwrap(); diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs index ae2b129ac8..7253f1559d 100644 --- a/compute_tools/src/extension_server.rs +++ b/compute_tools/src/extension_server.rs @@ -30,7 +30,10 @@ fn get_pg_version(pgbin: &str) -> String { if human_version.contains("v15") { return "v15".to_string(); } - "v14".to_string() + else if human_version.contains("v14") { + return "v14".to_string(); + } + panic!("Unsuported postgres version {human_version}"); } async fn download_helper( diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index e0840cd9a6..175e0ab705 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -666,7 +666,11 @@ class NeonEnvBuilder: enable_remote_extensions=enable_remote_extensions, ) elif remote_storage_kind == RemoteStorageKind.REAL_S3: - self.enable_real_s3_remote_storage(test_name=test_name, force_enable=force_enable) + self.enable_real_s3_remote_storage( + test_name=test_name, + force_enable=force_enable, + enable_remote_extensions=enable_remote_extensions, + ) else: raise RuntimeError(f"Unknown storage type: {remote_storage_kind}") @@ -722,7 +726,7 @@ class NeonEnvBuilder: secret_key=self.mock_s3_server.secret_key(), ) - def enable_real_s3_remote_storage(self, test_name: str, force_enable: bool = True): + def enable_real_s3_remote_storage(self, test_name: str, force_enable: bool = True, enable_remote_extensions: bool = False): """ Sets up configuration to use real s3 endpoint without mock server """ @@ -762,9 +766,10 @@ class NeonEnvBuilder: prefix_in_bucket=self.remote_storage_prefix, ) - ext_bucket_name = os.getenv("EXT_REMOTE_STORAGE_S3_BUCKET") - if ext_bucket_name is not None: - ext_bucket_name = f"ext_{ext_bucket_name}" + if enable_remote_extensions: + ext_bucket_name = os.getenv("EXT_REMOTE_STORAGE_S3_BUCKET") + if ext_bucket_name is None: + ext_bucket_name = "neon-dev-extensions" self.ext_remote_storage = S3Storage( bucket_name=ext_bucket_name, bucket_region=region, diff --git a/test_runner/fixtures/types.py b/test_runner/fixtures/types.py index 7d179cc7fb..ef88e09de4 100644 --- a/test_runner/fixtures/types.py +++ b/test_runner/fixtures/types.py @@ -89,6 +89,9 @@ class TenantId(Id): def __repr__(self) -> str: return f'`TenantId("{self.id.hex()}")' + def __str__(self) -> str: + return self.id.hex() + class TimelineId(Id): def __repr__(self) -> str: diff --git a/test_runner/regress/test_download_extensions.py b/test_runner/regress/test_download_extensions.py index e7e4cf5a65..6f99821c61 100644 --- a/test_runner/regress/test_download_extensions.py +++ b/test_runner/regress/test_download_extensions.py @@ -3,63 +3,94 @@ import os from contextlib import closing from io import BytesIO +import pytest from fixtures.log_helper import log from fixtures.neon_fixtures import ( NeonEnvBuilder, RemoteStorageKind, ) +""" +TODO: +- **add tests with real S3 storage** +- libs/remote_storage/src/s3_bucket.rs TODO // TODO: if bucket prefix is empty, + the folder is prefixed with a "/" I think. Is this desired? -def test_file_download(neon_env_builder: NeonEnvBuilder): +- Handle LIBRARY exttensions +- how to add env variable EXT_REMOTE_STORAGE_S3_BUCKET? +""" + + +def ext_contents(owner, i): + output = f"""# mock {owner} extension{i} +comment = 'This is a mock extension' +default_version = '1.0' +module_pathname = '$libdir/test_ext{i}' +relocatable = true""" + return output + + +@pytest.mark.parametrize( + "remote_storage_kind", + [RemoteStorageKind.LOCAL_FS, RemoteStorageKind.MOCK_S3, RemoteStorageKind.REAL_S3], +) +def test_file_download(neon_env_builder: NeonEnvBuilder, remote_storage_kind: RemoteStorageKind): """ Tests we can download a file First we set up the mock s3 bucket by uploading test_ext.control to the bucket Then, we download test_ext.control from the bucket to pg_install/v15/share/postgresql/extension/ Finally, we list available extensions and assert that test_ext is present """ + + if remote_storage_kind != RemoteStorageKind.MOCK_S3: + # skip these for now + return None + neon_env_builder.enable_remote_storage( - remote_storage_kind=RemoteStorageKind.MOCK_S3, + remote_storage_kind=remote_storage_kind, test_name="test_file_download", enable_remote_extensions=True, ) neon_env_builder.num_safekeepers = 3 env = neon_env_builder.init_start() + tenant_id, _ = env.neon_cli.create_tenant() + env.neon_cli.create_timeline("test_file_download", tenant_id=tenant_id) assert env.ext_remote_storage is not None assert env.remote_storage_client is not None - TEST_EXT_PATH = "v14/share/postgresql/extension/test_ext.control" + NUM_EXT = 5 + PUB_EXT_ROOT = "v14/share/postgresql/extension" BUCKET_PREFIX = "5314225671" # this is the build number + cleanup_files = [] - # 4. Upload test_ext.control file to the bucket - # In the non-mock version this is done by CI/CD + # Upload test_ext{i}.control files to the bucket + # Note: In real life this is done by CI/CD + for i in range(NUM_EXT): + # public extensions + public_ext = BytesIO(bytes(ext_contents("public", i), "utf-8")) + remote_name = f"{BUCKET_PREFIX}/{PUB_EXT_ROOT}/test_ext{i}.control" + local_name = f"pg_install/{PUB_EXT_ROOT}/test_ext{i}.control" + env.remote_storage_client.upload_fileobj( + public_ext, env.ext_remote_storage.bucket_name, remote_name + ) + cleanup_files.append(local_name) - test_ext_file = BytesIO( - b"""# mock extension -comment = 'This is a mock extension' -default_version = '1.0' -module_pathname = '$libdir/test_ext' -relocatable = true - """ - ) - env.remote_storage_client.upload_fileobj( - test_ext_file, - env.ext_remote_storage.bucket_name, - os.path.join(BUCKET_PREFIX, TEST_EXT_PATH), - ) + # private extensions + private_ext = BytesIO(bytes(ext_contents(str(tenant_id), i), "utf-8")) + remote_name = f"{BUCKET_PREFIX}/{str(tenant_id)}/private_ext{i}.control" + local_name = f"pg_install/{PUB_EXT_ROOT}/private_ext{i}.control" + env.remote_storage_client.upload_fileobj( + private_ext, env.ext_remote_storage.bucket_name, remote_name + ) + cleanup_files.append(local_name) - # 5. Download file from the bucket to correct local location - # Later this will be replaced by our rust code - # resp = env.remote_storage_client.get_object( - # Bucket=env.ext_remote_storage.bucket_name, Key=os.path.join(BUCKET_PREFIX, TEST_EXT_PATH) - # ) - # response = resp["Body"] - # fname = f"pg_install/{TEST_EXT_PATH}" - # with open(fname, "wb") as f: - # f.write(response.read()) - - tenant, _ = env.neon_cli.create_tenant() - env.neon_cli.create_timeline("test_file_download", tenant_id=tenant) + # Rust will then download the control files from the bucket + # our rust code should obtain the same result as the following: + # env.remote_storage_client.get_object( + # Bucket=env.ext_remote_storage.bucket_name, + # Key=os.path.join(BUCKET_PREFIX, PUB_EXT_PATHS[0]) + # )["Body"].read() remote_ext_config = json.dumps( { @@ -70,21 +101,35 @@ relocatable = true } ) - # 6. Start endpoint and ensure that test_ext is present in select * from pg_available_extensions endpoint = env.endpoints.create_start( - "test_file_download", tenant_id=tenant, remote_ext_config=remote_ext_config + "test_file_download", tenant_id=tenant_id, remote_ext_config=remote_ext_config ) with closing(endpoint.connect()) as conn: with conn.cursor() as cur: - # test query: insert some values and select them + # example query: insert some values and select them cur.execute("CREATE TABLE t(key int primary key, value text)") for i in range(100): cur.execute(f"insert into t values({i}, {2*i})") cur.execute("select * from t") log.info(cur.fetchall()) - # the real test query: check that test_ext is present + # Test query: check that test_ext0 was successfully downloaded cur.execute("SELECT * FROM pg_available_extensions") all_extensions = [x[0] for x in cur.fetchall()] log.info(all_extensions) - assert "test_ext" in all_extensions + for i in range(NUM_EXT): + assert f"test_ext{i}" in all_extensions + assert f"private_ext{i}" in all_extensions + + # TODO: can create extension actually install an extension? + # cur.execute("CREATE EXTENSION test_ext0") + # log.info("**" * 100) + # log.info(cur.fetchall()) + + # cleanup downloaded extensions (TODO: the file names are quesionable here) + for file in cleanup_files: + try: + log.info(f"Deleting {file}") + os.remove(file) + except FileNotFoundError: + log.info(f"{file} does not exist, so cannot be deleted")