rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-01-15 09:22:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

More Extension Features (#4555)

Browse Source 
Added tenant specific extensions and more tests
This commit is contained in:
Alek Westover
2023-06-23 09:30:49 -04:00
committed by GitHub
parent fd3dfe9d52
commit 31aa0283b0
5 changed files with 115 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
compute_tools/src/bin/compute_ctl.rs
View File

@@ -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(&copy_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();

5
compute_tools/src/extension_server.rs
View File

@@ -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(

15
test_runner/fixtures/neon_fixtures.py
View File

@@ -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,

3
test_runner/fixtures/types.py
View File

@@ -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:

115
test_runner/regress/test_download_extensions.py
View File

@@ -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")