test_runner: add --pg-version pytest argument (#4037)

- allows setting Postgres version for testing using --pg-version argument
- fixes tests for the non-default Postgres version.

test_runner/README.md

@@ -71,7 +71,8 @@ a subdirectory for each version with naming convention `v{PG_VERSION}/`.
 Inside that dir, a `bin/postgres` binary should be present.
 `DEFAULT_PG_VERSION`: The version of Postgres to use,
 This is used to construct full path to the postgres binaries.
-Format is 2-digit major version nubmer, i.e. `DEFAULT_PG_VERSION="14"`
+Format is 2-digit major version nubmer, i.e. `DEFAULT_PG_VERSION="14"`. Alternatively,
+you can use `--pg-version` argument.
 `TEST_OUTPUT`: Set the directory where test state and test output files
 should go.
 `TEST_SHARED_FIXTURES`: Try to re-use a single pageserver for all the tests.

test_runner/conftest.py

@@ -1,4 +1,5 @@
 pytest_plugins = (
+    "fixtures.pg_version",
     "fixtures.neon_fixtures",
     "fixtures.benchmark_fixture",
     "fixtures.pg_stats",

test_runner/fixtures/neon_fixtures.py

@@ -45,6 +45,7 @@ from typing_extensions import Literal
 from fixtures.log_helper import log
 from fixtures.pageserver.http import PageserverHttpClient
 from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload
+from fixtures.pg_version import PgVersion
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import (
     ATTACHMENT_NAME_REGEX,
@@ -75,7 +76,6 @@ Env = Dict[str, str]
 
 DEFAULT_OUTPUT_DIR: str = "test_output"
 DEFAULT_BRANCH_NAME: str = "main"
-DEFAULT_PG_VERSION_DEFAULT: str = "14"
 
 BASE_PORT: int = 15000
 WORKER_PORT_NUM: int = 1000
@@ -148,18 +148,7 @@ def top_output_dir(base_dir: Path) -> Iterator[Path]:
 
 
 @pytest.fixture(scope="session")
-def pg_version() -> Iterator[str]:
-    if env_default_pg_version := os.environ.get("DEFAULT_PG_VERSION"):
-        version = env_default_pg_version
-    else:
-        version = DEFAULT_PG_VERSION_DEFAULT
-
-    log.info(f"pg_version is {version}")
-    yield version
-
-
-@pytest.fixture(scope="session")
-def versioned_pg_distrib_dir(pg_distrib_dir: Path, pg_version: str) -> Iterator[Path]:
+def versioned_pg_distrib_dir(pg_distrib_dir: Path, pg_version: PgVersion) -> Iterator[Path]:
     versioned_dir = pg_distrib_dir / f"v{pg_version}"
 
     psql_bin_path = versioned_dir / "bin/psql"
@@ -592,7 +581,7 @@ class NeonEnvBuilder:
         mock_s3_server: MockS3Server,
         neon_binpath: Path,
         pg_distrib_dir: Path,
-        pg_version: str,
+        pg_version: PgVersion,
         remote_storage: Optional[RemoteStorage] = None,
         remote_storage_users: RemoteStorageUsers = RemoteStorageUsers.PAGESERVER,
         pageserver_config_override: Optional[str] = None,
@@ -1041,7 +1030,7 @@ def _shared_simple_env(
     top_output_dir: Path,
     neon_binpath: Path,
     pg_distrib_dir: Path,
-    pg_version: str,
+    pg_version: PgVersion,
 ) -> Iterator[NeonEnv]:
     """
     # Internal fixture backing the `neon_simple_env` fixture. If TEST_SHARED_FIXTURES
@@ -1098,7 +1087,7 @@ def neon_env_builder(
     mock_s3_server: MockS3Server,
     neon_binpath: Path,
     pg_distrib_dir: Path,
-    pg_version: str,
+    pg_version: PgVersion,
     default_broker: NeonBroker,
     run_id: uuid.UUID,
 ) -> Iterator[NeonEnvBuilder]:
@@ -1753,7 +1742,7 @@ def append_pageserver_param_overrides(
 class PgBin:
     """A helper class for executing postgres binaries"""
 
-    def __init__(self, log_dir: Path, pg_distrib_dir: Path, pg_version: str):
+    def __init__(self, log_dir: Path, pg_distrib_dir: Path, pg_version: PgVersion):
         self.log_dir = log_dir
         self.pg_version = pg_version
         self.pg_bin_path = pg_distrib_dir / f"v{pg_version}" / "bin"
@@ -1812,7 +1801,7 @@ class PgBin:
 
 
 @pytest.fixture(scope="function")
-def pg_bin(test_output_dir: Path, pg_distrib_dir: Path, pg_version: str) -> PgBin:
+def pg_bin(test_output_dir: Path, pg_distrib_dir: Path, pg_version: PgVersion) -> PgBin:
     return PgBin(test_output_dir, pg_distrib_dir, pg_version)
 
 
@@ -1900,7 +1889,7 @@ def vanilla_pg(
     test_output_dir: Path,
     port_distributor: PortDistributor,
     pg_distrib_dir: Path,
-    pg_version: str,
+    pg_version: PgVersion,
 ) -> Iterator[VanillaPostgres]:
     pgdatadir = test_output_dir / "pgdata-vanilla"
     pg_bin = PgBin(test_output_dir, pg_distrib_dir, pg_version)
@@ -1945,7 +1934,7 @@ class RemotePostgres(PgProtocol):
 
 @pytest.fixture(scope="function")
 def remote_pg(
-    test_output_dir: Path, pg_distrib_dir: Path, pg_version: str
+    test_output_dir: Path, pg_distrib_dir: Path, pg_version: PgVersion
 ) -> Iterator[RemotePostgres]:
     pg_bin = PgBin(test_output_dir, pg_distrib_dir, pg_version)
 
@@ -2629,7 +2618,7 @@ class Safekeeper:
 @dataclass
 class SafekeeperTimelineStatus:
     acceptor_epoch: int
-    pg_version: int
+    pg_version: int  # Not exactly a PgVersion, safekeeper returns version as int, for example 150002 for 15.2
     flush_lsn: Lsn
     commit_lsn: Lsn
     timeline_start_lsn: Lsn
@@ -2675,7 +2664,11 @@ class SafekeeperHttpClient(requests.Session):
         return res_json
 
     def timeline_create(
-        self, tenant_id: TenantId, timeline_id: TimelineId, pg_version: int, commit_lsn: Lsn
+        self,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        pg_version: int,  # Not exactly a PgVersion, safekeeper returns version as int, for example 150002 for 15.2
+        commit_lsn: Lsn,
     ):
         body = {
             "tenant_id": str(tenant_id),

test_runner/fixtures/pageserver/http.py

@@ -9,6 +9,7 @@ import requests
 
 from fixtures.log_helper import log
 from fixtures.metrics import Metrics, parse_metrics
+from fixtures.pg_version import PgVersion
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import Fn
 
@@ -266,19 +267,21 @@ class PageserverHttpClient(requests.Session):
 
     def timeline_create(
         self,
+        pg_version: PgVersion,
         tenant_id: TenantId,
         new_timeline_id: Optional[TimelineId] = None,
         ancestor_timeline_id: Optional[TimelineId] = None,
         ancestor_start_lsn: Optional[Lsn] = None,
     ) -> Dict[Any, Any]:
-        res = self.post(
-            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline",
-            json={
-                "new_timeline_id": str(new_timeline_id) if new_timeline_id else None,
-                "ancestor_start_lsn": str(ancestor_start_lsn) if ancestor_start_lsn else None,
-                "ancestor_timeline_id": str(ancestor_timeline_id) if ancestor_timeline_id else None,
-            },
-        )
+        body: Dict[str, Any] = {
+            "new_timeline_id": str(new_timeline_id) if new_timeline_id else None,
+            "ancestor_start_lsn": str(ancestor_start_lsn) if ancestor_start_lsn else None,
+            "ancestor_timeline_id": str(ancestor_timeline_id) if ancestor_timeline_id else None,
+        }
+        if pg_version != PgVersion.NOT_SET:
+            body["pg_version"] = int(pg_version)
+
+        res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline", json=body)
         self.verbose_error(res)
         if res.status_code == 409:
             raise Exception(f"could not create timeline: already exists for id {new_timeline_id}")

test_runner/fixtures/pg_version.py

@@ -0,0 +1,68 @@
+import enum
+import os
+from typing import Iterator, Optional
+
+import pytest
+from _pytest.config.argparsing import Parser
+from pytest import FixtureRequest
+
+from fixtures.log_helper import log
+
+"""
+This fixture is used to determine which version of Postgres to use for tests.
+"""
+
+
+# Inherit PgVersion from str rather than int to make it easier to pass as a command-line argument
+# TODO: use enum.StrEnum for Python >= 3.11
+@enum.unique
+class PgVersion(str, enum.Enum):
+    V14 = "14"
+    V15 = "15"
+    # Instead of making version an optional parameter in methods, we can use this fake entry
+    # to explicitly rely on the default server version (could be different from pg_version fixture value)
+    NOT_SET = "<-POSTRGRES VERSION IS NOT SET->"
+
+    # Make it less confusing in logs
+    def __repr__(self) -> str:
+        return f"'{self.value}'"
+
+    @classmethod
+    def _missing_(cls, value) -> Optional["PgVersion"]:
+        known_values = {v.value for _, v in cls.__members__.items()}
+
+        # Allow passing version as a string with "v" prefix (e.g. "v14")
+        if isinstance(value, str) and value.lower().startswith("v") and value[1:] in known_values:
+            return cls(value[1:])
+        # Allow passing version as an int (e.g. 15 or 150002, both will be converted to PgVersion.V15)
+        elif isinstance(value, int) and str(value)[:2] in known_values:
+            return cls(str(value)[:2])
+
+        # Make mypy happy
+        # See https://github.com/python/mypy/issues/3974
+        return None
+
+
+DEFAULT_VERSION: PgVersion = PgVersion.V14
+
+
+def pytest_addoption(parser: Parser):
+    parser.addoption(
+        "--pg-version",
+        action="store",
+        type=PgVersion,
+        help="Postgres version to use for tests",
+    )
+
+
+@pytest.fixture(scope="session")
+def pg_version(request: FixtureRequest) -> Iterator[PgVersion]:
+    if v := request.config.getoption("--pg-version"):
+        version, source = v, "from --pg-version commad-line argument"
+    elif v := os.environ.get("DEFAULT_PG_VERSION"):
+        version, source = PgVersion(v), "from DEFAULT_PG_VERSION environment variable"
+    else:
+        version, source = DEFAULT_VERSION, "default verson"
+
+    log.info(f"pg_version is {version} ({source})")
+    yield version

test_runner/regress/test_auth.py

@@ -31,11 +31,15 @@ def test_pageserver_auth(neon_env_builder: NeonEnvBuilder):
 
     # tenant can create branches
     tenant_http_client.timeline_create(
-        tenant_id=env.initial_tenant, ancestor_timeline_id=new_timeline_id
+        pg_version=env.pg_version,
+        tenant_id=env.initial_tenant,
+        ancestor_timeline_id=new_timeline_id,
     )
     # console can create branches for tenant
     pageserver_http_client.timeline_create(
-        tenant_id=env.initial_tenant, ancestor_timeline_id=new_timeline_id
+        pg_version=env.pg_version,
+        tenant_id=env.initial_tenant,
+        ancestor_timeline_id=new_timeline_id,
     )
 
     # fail to create branch using token with different tenant_id
@@ -43,7 +47,9 @@ def test_pageserver_auth(neon_env_builder: NeonEnvBuilder):
         PageserverApiException, match="Forbidden: Tenant id mismatch. Permission denied"
     ):
         invalid_tenant_http_client.timeline_create(
-            tenant_id=env.initial_tenant, ancestor_timeline_id=new_timeline_id
+            pg_version=env.pg_version,
+            tenant_id=env.initial_tenant,
+            ancestor_timeline_id=new_timeline_id,
         )
 
     # create tenant using management token

test_runner/regress/test_compatibility.py

@@ -16,6 +16,7 @@ from fixtures.neon_fixtures import (
 )
 from fixtures.pageserver.http import PageserverHttpClient
 from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload
+from fixtures.pg_version import PgVersion
 from fixtures.types import Lsn
 from pytest import FixtureRequest
 
@@ -50,7 +51,7 @@ def test_create_snapshot(neon_env_builder: NeonEnvBuilder, pg_bin: PgBin, test_o
     # it creates a new snapshot for releases after we tested the current version against the previous snapshot in `test_backward_compatibility`.
     #
     # There's no cleanup here, it allows to adjust the data in `test_backward_compatibility` itself without re-collecting it.
-    neon_env_builder.pg_version = "14"
+    neon_env_builder.pg_version = PgVersion.V14
     neon_env_builder.num_safekeepers = 3
     neon_env_builder.enable_local_fs_remote_storage()
     neon_env_builder.preserve_database_files = True
@@ -98,7 +99,7 @@ def test_backward_compatibility(
     test_output_dir: Path,
     neon_binpath: Path,
     pg_distrib_dir: Path,
-    pg_version: str,
+    pg_version: PgVersion,
     request: FixtureRequest,
 ):
     """
@@ -153,7 +154,7 @@ def test_backward_compatibility(
 def test_forward_compatibility(
     test_output_dir: Path,
     port_distributor: PortDistributor,
-    pg_version: str,
+    pg_version: PgVersion,
     request: FixtureRequest,
     neon_binpath: Path,
 ):
@@ -345,7 +346,7 @@ def check_neon_works(
     neon_target_binpath: Path,
     neon_current_binpath: Path,
     pg_distrib_dir: Path,
-    pg_version: str,
+    pg_version: PgVersion,
     port_distributor: PortDistributor,
     test_output_dir: Path,
     pg_bin: PgBin,
@@ -404,7 +405,7 @@ def check_neon_works(
 
     shutil.rmtree(repo_dir / "local_fs_remote_storage")
     pageserver_http.timeline_delete(tenant_id, timeline_id)
-    pageserver_http.timeline_create(tenant_id, timeline_id)
+    pageserver_http.timeline_create(pg_version, tenant_id, timeline_id)
     pg_bin.run(
         ["pg_dumpall", f"--dbname={connstr}", f"--file={test_output_dir / 'dump-from-wal.sql'}"]
     )

test_runner/regress/test_pageserver_api.py

@@ -8,6 +8,7 @@ from fixtures.neon_fixtures import (
     NeonEnvBuilder,
 )
 from fixtures.pageserver.http import PageserverHttpClient
+from fixtures.pg_version import PgVersion
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import wait_until
 
@@ -61,7 +62,7 @@ def test_pageserver_init_node_id(
     assert "has node id already, it cannot be overridden" in bad_update.stderr
 
 
-def check_client(client: PageserverHttpClient, initial_tenant: TenantId):
+def check_client(pg_version: PgVersion, client: PageserverHttpClient, initial_tenant: TenantId):
     client.check_status()
 
     # check initial tenant is there
@@ -77,7 +78,11 @@ def check_client(client: PageserverHttpClient, initial_tenant: TenantId):
 
     # create timeline
     timeline_id = TimelineId.generate()
-    client.timeline_create(tenant_id=tenant_id, new_timeline_id=timeline_id)
+    client.timeline_create(
+        pg_version=pg_version,
+        tenant_id=tenant_id,
+        new_timeline_id=timeline_id,
+    )
 
     timelines = client.timeline_list(tenant_id)
     assert len(timelines) > 0
@@ -174,7 +179,7 @@ def test_pageserver_http_get_wal_receiver_success(neon_simple_env: NeonEnv):
 def test_pageserver_http_api_client(neon_simple_env: NeonEnv):
     env = neon_simple_env
     with env.pageserver.http_client() as client:
-        check_client(client, env.initial_tenant)
+        check_client(env.pg_version, client, env.initial_tenant)
 
 
 def test_pageserver_http_api_client_auth_enabled(neon_env_builder: NeonEnvBuilder):
@@ -184,4 +189,4 @@ def test_pageserver_http_api_client_auth_enabled(neon_env_builder: NeonEnvBuilde
     pageserver_token = env.auth_keys.generate_pageserver_token()
 
     with env.pageserver.http_client(auth_token=pageserver_token) as client:
-        check_client(client, env.initial_tenant)
+        check_client(env.pg_version, client, env.initial_tenant)

test_runner/regress/test_timeline_size.py

@@ -27,6 +27,7 @@ from fixtures.pageserver.utils import (
     wait_for_upload_queue_empty,
     wait_until_tenant_active,
 )
+from fixtures.pg_version import PgVersion
 from fixtures.types import TenantId, TimelineId
 from fixtures.utils import get_timeline_dir_size, wait_until
 
@@ -489,7 +490,7 @@ def test_timeline_size_metrics(
     test_output_dir: Path,
     port_distributor: PortDistributor,
     pg_distrib_dir: Path,
-    pg_version: str,
+    pg_version: PgVersion,
 ):
     env = neon_simple_env
     pageserver_http = env.pageserver.http_client()

test_runner/regress/test_wal_acceptor.py

@@ -32,6 +32,7 @@ from fixtures.neon_fixtures import (
     available_remote_storages,
 )
 from fixtures.pageserver.utils import wait_for_last_record_lsn, wait_for_upload
+from fixtures.pg_version import PgVersion
 from fixtures.types import Lsn, TenantId, TimelineId
 from fixtures.utils import get_dir_size, query_scalar, start_in_background
 
@@ -582,7 +583,11 @@ def test_s3_wal_replay(neon_env_builder: NeonEnvBuilder, remote_storage_kind: Re
         shutil.copy(f_partial_saved, f_partial_path)
 
     # recreate timeline on pageserver from scratch
-    ps_cli.timeline_create(tenant_id, timeline_id)
+    ps_cli.timeline_create(
+        pg_version=PgVersion(pg_version),
+        tenant_id=tenant_id,
+        new_timeline_id=timeline_id,
+    )
 
     wait_lsn_timeout = 60 * 3
     started_at = time.time()