diff --git a/test_runner/fixtures/zenith_fixtures.py b/test_runner/fixtures/zenith_fixtures.py index f581c8a3ee..f7665c89d8 100644 --- a/test_runner/fixtures/zenith_fixtures.py +++ b/test_runner/fixtures/zenith_fixtures.py @@ -25,7 +25,7 @@ from dataclasses import dataclass # Type-related stuff from psycopg2.extensions import connection as PgConnection -from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar, cast +from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar, cast, Union from typing_extensions import Literal import pytest @@ -342,10 +342,12 @@ class ZenithEnvBuilder: def __init__(self, repo_dir: Path, port_distributor: PortDistributor, + pageserver_remote_storage: Optional[RemoteStorage] = None, num_safekeepers: int = 0, pageserver_auth_enabled: bool = False): self.repo_dir = repo_dir self.port_distributor = port_distributor + self.pageserver_remote_storage = pageserver_remote_storage self.num_safekeepers = num_safekeepers self.pageserver_auth_enabled = pageserver_auth_enabled self.env: Optional[ZenithEnv] = None @@ -434,7 +436,9 @@ auth_type = '{pageserver_auth_type}' """ # Create a corresponding ZenithPageserver object - self.pageserver = ZenithPageserver(self, port=pageserver_port) + self.pageserver = ZenithPageserver(self, + port=pageserver_port, + remote_storage=config.pageserver_remote_storage) # Create config and a Safekeeper object for each safekeeper for i in range(1, config.num_safekeepers + 1): @@ -465,6 +469,8 @@ sync = false # Disable fsyncs to make the tests go faster tmp.flush() cmd = ['init', f'--config={tmp.name}'] + _append_pageserver_param_overrides(cmd, config.pageserver_remote_storage) + self.zenith_cli(cmd) # Start up the page server and all the safekeepers @@ -677,17 +683,38 @@ class PageserverPort: http: int +@dataclass +class LocalFsStorage: + root: Path + + +@dataclass +class S3Storage: + bucket: str + region: str + access_key: str + secret_key: str + + +RemoteStorage = Union[LocalFsStorage, S3Storage] + + class ZenithPageserver(PgProtocol): """ An object representing a running pageserver. Initializes the repository via `zenith init`. """ - def __init__(self, env: ZenithEnv, port: PageserverPort, enable_auth=False): + def __init__(self, + env: ZenithEnv, + port: PageserverPort, + remote_storage: Optional[RemoteStorage] = None, + enable_auth=False): super().__init__(host='localhost', port=port.pg) self.env = env self.running = False self.service_port = port # do not shadow PgProtocol.port which is just int + self.remote_storage = remote_storage def start(self) -> 'ZenithPageserver': """ @@ -696,7 +723,10 @@ class ZenithPageserver(PgProtocol): """ assert self.running == False - self.env.zenith_cli(['pageserver', 'start']) + start_args = ['pageserver', 'start'] + _append_pageserver_param_overrides(start_args, self.remote_storage) + + self.env.zenith_cli(start_args) self.running = True return self @@ -729,6 +759,28 @@ class ZenithPageserver(PgProtocol): ) +def _append_pageserver_param_overrides(params_to_update: List[str], + pageserver_remote_storage: Optional[RemoteStorage]): + if pageserver_remote_storage is not None: + if isinstance(pageserver_remote_storage, LocalFsStorage): + pageserver_storage_override = f"local_path='{pageserver_remote_storage.root}'" + elif isinstance(pageserver_remote_storage, S3Storage): + pageserver_storage_override = f"bucket_name='{pageserver_remote_storage.bucket}',\ + bucket_region='{pageserver_remote_storage.region}',access_key_id='{pageserver_remote_storage.access_key}',\ + secret_access_key='{pageserver_remote_storage.secret_key}'" + + else: + raise Exception(f'Unknown storage configuration {pageserver_remote_storage}') + params_to_update.append( + f'--pageserver-config-override=remote_storage={{{pageserver_storage_override}}}') + + env_overrides = os.getenv('ZENITH_PAGESERVER_OVERRIDES') + if env_overrides is not None: + params_to_update += [ + f'--pageserver-config-override={o.strip()}' for o in env_overrides.split(';') + ] + + class PgBin: """ A helper class for executing postgres binaries """ def __init__(self, log_dir: str):