Support custom extensions.

Add infrastructure to dynamically load postgres extensions and shared libraries from remote extension storage.

Before postgres start  downloads list of available remote extensions and libraries, and also downloads 'shared_preload_libraries'.
After postgres is running, 'compute_ctl' listens for HTTP requests to load files.

Postgres has new GUC 'extension_server_port' to specify port on which 'compute_ctl' listens for requests.

When PostgreSQL requests a file, 'compute_ctl'  downloads it.

See more details about feature design and remote extension storage layout in docs/rfcs/024-extension-loading.md

test_runner/fixtures/neon_fixtures.py

@@ -514,6 +514,16 @@ def available_remote_storages() -> List[RemoteStorageKind]:
     return remote_storages
 
 
+def available_s3_storages() -> List[RemoteStorageKind]:
+    remote_storages = [RemoteStorageKind.MOCK_S3]
+    if os.getenv("ENABLE_REAL_S3_REMOTE_STORAGE") is not None:
+        remote_storages.append(RemoteStorageKind.REAL_S3)
+        log.info("Enabling real s3 storage for tests")
+    else:
+        log.info("Using mock implementations to test remote storage")
+    return remote_storages
+
+
 @dataclass
 class LocalFsStorage:
     root: Path
@@ -534,6 +544,16 @@ class S3Storage:
             "AWS_SECRET_ACCESS_KEY": self.secret_key,
         }
 
+    def to_string(self) -> str:
+        return json.dumps(
+            {
+                "bucket": self.bucket_name,
+                "region": self.bucket_region,
+                "endpoint": self.endpoint,
+                "prefix": self.prefix_in_bucket,
+            }
+        )
+
 
 RemoteStorage = Union[LocalFsStorage, S3Storage]
 
@@ -600,10 +620,12 @@ class NeonEnvBuilder:
         self.rust_log_override = rust_log_override
         self.port_distributor = port_distributor
         self.remote_storage = remote_storage
+        self.ext_remote_storage: Optional[S3Storage] = None
+        self.remote_storage_client: Optional[Any] = None
         self.remote_storage_users = remote_storage_users
         self.broker = broker
         self.run_id = run_id
-        self.mock_s3_server = mock_s3_server
+        self.mock_s3_server: MockS3Server = mock_s3_server
         self.pageserver_config_override = pageserver_config_override
         self.num_safekeepers = num_safekeepers
         self.safekeepers_id_start = safekeepers_id_start
@@ -651,15 +673,24 @@ class NeonEnvBuilder:
         remote_storage_kind: RemoteStorageKind,
         test_name: str,
         force_enable: bool = True,
+        enable_remote_extensions: bool = False,
     ):
         if remote_storage_kind == RemoteStorageKind.NOOP:
             return
         elif remote_storage_kind == RemoteStorageKind.LOCAL_FS:
             self.enable_local_fs_remote_storage(force_enable=force_enable)
         elif remote_storage_kind == RemoteStorageKind.MOCK_S3:
-            self.enable_mock_s3_remote_storage(bucket_name=test_name, force_enable=force_enable)
+            self.enable_mock_s3_remote_storage(
+                bucket_name=test_name,
+                force_enable=force_enable,
+                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}")
 
@@ -673,11 +704,15 @@ class NeonEnvBuilder:
         assert force_enable or self.remote_storage is None, "remote storage is enabled already"
         self.remote_storage = LocalFsStorage(Path(self.repo_dir / "local_fs_remote_storage"))
 
-    def enable_mock_s3_remote_storage(self, bucket_name: str, force_enable: bool = True):
+    def enable_mock_s3_remote_storage(
+        self, bucket_name: str, force_enable: bool = True, enable_remote_extensions: bool = False
+    ):
         """
         Sets up the pageserver to use the S3 mock server, creates the bucket, if it's not present already.
         Starts up the mock server, if that does not run yet.
         Errors, if the pageserver has some remote storage configuration already, unless `force_enable` is not set to `True`.
+
+        Also creates the bucket for extensions, self.ext_remote_storage bucket
         """
         assert force_enable or self.remote_storage is None, "remote storage is enabled already"
         mock_endpoint = self.mock_s3_server.endpoint()
@@ -698,9 +733,22 @@ class NeonEnvBuilder:
             bucket_region=mock_region,
             access_key=self.mock_s3_server.access_key(),
             secret_key=self.mock_s3_server.secret_key(),
+            prefix_in_bucket="pageserver",
         )
 
-    def enable_real_s3_remote_storage(self, test_name: str, force_enable: bool = True):
+        if enable_remote_extensions:
+            self.ext_remote_storage = S3Storage(
+                bucket_name=bucket_name,
+                endpoint=mock_endpoint,
+                bucket_region=mock_region,
+                access_key=self.mock_s3_server.access_key(),
+                secret_key=self.mock_s3_server.secret_key(),
+                prefix_in_bucket="ext",
+            )
+
+    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
         """
@@ -737,9 +785,18 @@ class NeonEnvBuilder:
             bucket_region=region,
             access_key=access_key,
             secret_key=secret_key,
-            prefix_in_bucket=self.remote_storage_prefix,
+            prefix_in_bucket=f"{self.remote_storage_prefix}/pageserver",
         )
 
+        if enable_remote_extensions:
+            self.ext_remote_storage = S3Storage(
+                bucket_name=bucket_name,
+                bucket_region=region,
+                access_key=access_key,
+                secret_key=secret_key,
+                prefix_in_bucket=f"{self.remote_storage_prefix}/ext",
+            )
+
     def cleanup_local_storage(self):
         if self.preserve_database_files:
             return
@@ -773,6 +830,7 @@ class NeonEnvBuilder:
         # `self.remote_storage_prefix` is coupled with `S3Storage` storage type,
         # so this line effectively a no-op
         assert isinstance(self.remote_storage, S3Storage)
+        assert self.remote_storage_client is not None
 
         if self.keep_remote_storage_contents:
             log.info("keep_remote_storage_contents skipping remote storage cleanup")
@@ -902,6 +960,8 @@ class NeonEnv:
         self.neon_binpath = config.neon_binpath
         self.pg_distrib_dir = config.pg_distrib_dir
         self.endpoint_counter = 0
+        self.remote_storage_client = config.remote_storage_client
+        self.ext_remote_storage = config.ext_remote_storage
 
         # generate initial tenant ID here instead of letting 'neon init' generate it,
         # so that we don't need to dig it out of the config file afterwards.
@@ -1488,6 +1548,7 @@ class NeonCli(AbstractNeonCli):
         safekeepers: Optional[List[int]] = None,
         tenant_id: Optional[TenantId] = None,
         lsn: Optional[Lsn] = None,
+        remote_ext_config: Optional[str] = None,
     ) -> "subprocess.CompletedProcess[str]":
         args = [
             "endpoint",
@@ -1497,6 +1558,8 @@ class NeonCli(AbstractNeonCli):
             "--pg-version",
             self.env.pg_version,
         ]
+        if remote_ext_config is not None:
+            args.extend(["--remote-ext-config", remote_ext_config])
         if lsn is not None:
             args.append(f"--lsn={lsn}")
         args.extend(["--pg-port", str(pg_port)])
@@ -2358,7 +2421,7 @@ class Endpoint(PgProtocol):
 
         return self
 
-    def start(self) -> "Endpoint":
+    def start(self, remote_ext_config: Optional[str] = None) -> "Endpoint":
         """
         Start the Postgres instance.
         Returns self.
@@ -2374,6 +2437,7 @@ class Endpoint(PgProtocol):
             http_port=self.http_port,
             tenant_id=self.tenant_id,
             safekeepers=self.active_safekeepers,
+            remote_ext_config=remote_ext_config,
         )
         self.running = True
 
@@ -2463,6 +2527,7 @@ class Endpoint(PgProtocol):
         hot_standby: bool = False,
         lsn: Optional[Lsn] = None,
         config_lines: Optional[List[str]] = None,
+        remote_ext_config: Optional[str] = None,
     ) -> "Endpoint":
         """
         Create an endpoint, apply config, and start Postgres.
@@ -2477,7 +2542,7 @@ class Endpoint(PgProtocol):
             config_lines=config_lines,
             hot_standby=hot_standby,
             lsn=lsn,
-        ).start()
+        ).start(remote_ext_config=remote_ext_config)
 
         log.info(f"Postgres startup took {time.time() - started_at} seconds")
 
@@ -2511,6 +2576,7 @@ class EndpointFactory:
         lsn: Optional[Lsn] = None,
         hot_standby: bool = False,
         config_lines: Optional[List[str]] = None,
+        remote_ext_config: Optional[str] = None,
     ) -> Endpoint:
         ep = Endpoint(
             self.env,
@@ -2527,6 +2593,7 @@ class EndpointFactory:
             hot_standby=hot_standby,
             config_lines=config_lines,
             lsn=lsn,
+            remote_ext_config=remote_ext_config,
         )
 
     def create(

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:

test_runner/regress/test_download_extensions.py

@@ -0,0 +1,247 @@
+import os
+from contextlib import closing
+from io import BytesIO
+
+import pytest
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import (
+    NeonEnvBuilder,
+    PgBin,
+    RemoteStorageKind,
+    available_s3_storages,
+)
+from fixtures.pg_version import PgVersion
+from fixtures.types import TenantId
+
+NUM_EXT = 3
+
+
+def control_file_content(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
+
+
+def sql_file_content():
+    output = """
+            CREATE FUNCTION test_ext_add(integer, integer) RETURNS integer
+    AS 'select $1 + $2;'
+    LANGUAGE SQL
+    IMMUTABLE
+    RETURNS NULL ON NULL INPUT;
+        """
+    return output
+
+
+# Prepare some mock extension files and upload them to the bucket
+# returns a list of files that should be cleaned up after the test
+def prepare_mock_ext_storage(
+    pg_version: PgVersion,
+    tenant_id: TenantId,
+    pg_bin: PgBin,
+    ext_remote_storage,
+    remote_storage_client,
+):
+    bucket_prefix = ext_remote_storage.prefix_in_bucket
+    custom_prefix = str(tenant_id)
+    PUB_EXT_ROOT = f"v{pg_version}/share/postgresql/extension"
+    PRIVATE_EXT_ROOT = f"v{pg_version}/{custom_prefix}/share/postgresql/extension"
+    LOCAL_EXT_ROOT = f"pg_install/{PUB_EXT_ROOT}"
+
+    PUB_LIB_ROOT = f"v{pg_version}/lib"
+    PRIVATE_LIB_ROOT = f"v{pg_version}/{custom_prefix}/lib"
+    LOCAL_LIB_ROOT = f"{pg_bin.pg_lib_dir}/postgresql"
+
+    log.info(
+        f"""
+            PUB_EXT_ROOT: {PUB_EXT_ROOT}
+            PRIVATE_EXT_ROOT: {PRIVATE_EXT_ROOT}
+            LOCAL_EXT_ROOT: {LOCAL_EXT_ROOT}
+            PUB_LIB_ROOT: {PUB_LIB_ROOT}
+            PRIVATE_LIB_ROOT: {PRIVATE_LIB_ROOT}
+            LOCAL_LIB_ROOT: {LOCAL_LIB_ROOT}
+            """
+    )
+
+    cleanup_files = []
+
+    # Upload several test_ext{i}.control files to the bucket
+    for i in range(NUM_EXT):
+        public_ext = BytesIO(bytes(control_file_content("public", i), "utf-8"))
+        public_remote_name = f"{bucket_prefix}/{PUB_EXT_ROOT}/test_ext{i}.control"
+        public_local_name = f"{LOCAL_EXT_ROOT}/test_ext{i}.control"
+        custom_ext = BytesIO(bytes(control_file_content(str(tenant_id), i), "utf-8"))
+        custom_remote_name = f"{bucket_prefix}/{PRIVATE_EXT_ROOT}/custom_ext{i}.control"
+        custom_local_name = f"{LOCAL_EXT_ROOT}/custom_ext{i}.control"
+        cleanup_files += [public_local_name, custom_local_name]
+
+        remote_storage_client.upload_fileobj(
+            public_ext, ext_remote_storage.bucket_name, public_remote_name
+        )
+        remote_storage_client.upload_fileobj(
+            custom_ext, ext_remote_storage.bucket_name, custom_remote_name
+        )
+
+    # Upload SQL file for the extension we're going to create
+    sql_filename = "test_ext0--1.0.sql"
+    test_sql_public_remote_path = f"{bucket_prefix}/{PUB_EXT_ROOT}/{sql_filename}"
+    test_sql_local_path = f"{LOCAL_EXT_ROOT}/{sql_filename}"
+    test_ext_sql_file = BytesIO(bytes(sql_file_content(), "utf-8"))
+    remote_storage_client.upload_fileobj(
+        test_ext_sql_file,
+        ext_remote_storage.bucket_name,
+        test_sql_public_remote_path,
+    )
+    cleanup_files += [test_sql_local_path]
+
+    # upload some fake library files
+    for i in range(2):
+        public_library = BytesIO(bytes("\n111\n", "utf-8"))
+        public_remote_name = f"{bucket_prefix}/{PUB_LIB_ROOT}/test_lib{i}.so"
+        public_local_name = f"{LOCAL_LIB_ROOT}/test_lib{i}.so"
+        custom_library = BytesIO(bytes("\n111\n", "utf-8"))
+        custom_remote_name = f"{bucket_prefix}/{PRIVATE_LIB_ROOT}/custom_lib{i}.so"
+        custom_local_name = f"{LOCAL_LIB_ROOT}/custom_lib{i}.so"
+
+        log.info(f"uploading library to {public_remote_name}")
+        log.info(f"uploading library to {custom_remote_name}")
+
+        remote_storage_client.upload_fileobj(
+            public_library,
+            ext_remote_storage.bucket_name,
+            public_remote_name,
+        )
+        remote_storage_client.upload_fileobj(
+            custom_library,
+            ext_remote_storage.bucket_name,
+            custom_remote_name,
+        )
+        cleanup_files += [public_local_name, custom_local_name]
+
+    return cleanup_files
+
+
+# Generate mock extension files and upload them to the bucket.
+#
+# Then check that compute nodes can download them and use them
+# to CREATE EXTENSION and LOAD 'library.so'
+#
+# NOTE: You must have appropriate AWS credentials to run REAL_S3 test.
+# It may also be necessary to set the following environment variables:
+#   export AWS_ACCESS_KEY_ID='test'
+#   export AWS_SECRET_ACCESS_KEY='test'
+#   export AWS_SECURITY_TOKEN='test'
+#   export AWS_SESSION_TOKEN='test'
+#   export AWS_DEFAULT_REGION='us-east-1'
+#
+@pytest.mark.parametrize("remote_storage_kind", available_s3_storages())
+def test_remote_extensions(
+    neon_env_builder: NeonEnvBuilder,
+    remote_storage_kind: RemoteStorageKind,
+    pg_version: PgVersion,
+    pg_bin: PgBin,
+):
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=remote_storage_kind,
+        test_name="test_remote_extensions",
+        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_remote_extensions", tenant_id=tenant_id)
+
+    assert env.ext_remote_storage is not None
+    assert env.remote_storage_client is not None
+
+    # Prepare some mock extension files and upload them to the bucket
+    cleanup_files = prepare_mock_ext_storage(
+        pg_version,
+        tenant_id,
+        pg_bin,
+        env.ext_remote_storage,
+        env.remote_storage_client,
+    )
+    # Start a compute node and check that it can download the extensions
+    # and use them to CREATE EXTENSION and LOAD 'library.so'
+    #
+    # This block is wrapped in a try/finally so that the downloaded files
+    # are cleaned up even if the test fails
+    try:
+        endpoint = env.endpoints.create_start(
+            "test_remote_extensions",
+            tenant_id=tenant_id,
+            remote_ext_config=env.ext_remote_storage.to_string(),
+            # config_lines=["log_min_messages=debug3"],
+        )
+        with closing(endpoint.connect()) as conn:
+            with conn.cursor() as cur:
+                # 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)
+                for i in range(NUM_EXT):
+                    assert f"test_ext{i}" in all_extensions
+                    assert f"custom_ext{i}" in all_extensions
+
+                cur.execute("CREATE EXTENSION test_ext0")
+                cur.execute("SELECT extname FROM pg_extension")
+                all_extensions = [x[0] for x in cur.fetchall()]
+                log.info(all_extensions)
+                assert "test_ext0" in all_extensions
+
+                # Try to load existing library file
+                try:
+                    cur.execute("LOAD 'test_lib0.so'")
+                except Exception as e:
+                    # expected to fail with
+                    # could not load library ... test_ext.so: file too short
+                    # because test_lib0.so is not real library file
+                    log.info("LOAD test_lib0.so failed (expectedly): %s", e)
+                    assert "file too short" in str(e)
+
+                # Try to load custom library file
+                try:
+                    cur.execute("LOAD 'custom_lib0.so'")
+                except Exception as e:
+                    # expected to fail with
+                    # could not load library ... test_ext.so: file too short
+                    # because test_lib0.so is not real library file
+                    log.info("LOAD custom_lib0.so failed (expectedly): %s", e)
+                    assert "file too short" in str(e)
+
+                # Try to load existing library file without .so extension
+                try:
+                    cur.execute("LOAD 'test_lib1'")
+                except Exception as e:
+                    # expected to fail with
+                    # could not load library ... test_lib1.so: file too short
+                    # because test_lib1.so is not real library file
+                    log.info("LOAD test_lib1 failed (expectedly): %s", e)
+                    assert "file too short" in str(e)
+
+                # Try to load non-existent library file
+                try:
+                    cur.execute("LOAD 'test_lib_fail.so'")
+                except Exception as e:
+                    # expected to fail because test_lib_fail.so is not found
+                    log.info("LOAD test_lib_fail.so failed (expectedly): %s", e)
+                    assert (
+                        """could not access file "test_lib_fail.so": No such file or directory"""
+                        in str(e)
+                    )
+
+    finally:
+        # this is important because if the files aren't cleaned up then the test can
+        # pass even without successfully downloading the files if a previous run (or
+        # run with different type of remote storage) of the test did download the
+        # files
+        for file in cleanup_files:
+            try:
+                os.remove(file)
+                log.info(f"Deleted {file}")
+            except FileNotFoundError:
+                log.info(f"{file} does not exist, so cannot be deleted")

test_runner/regress/test_timeline_delete.py

@@ -275,6 +275,7 @@ def assert_prefix_empty(neon_env_builder: NeonEnvBuilder, prefix: Optional[str]
     assert isinstance(neon_env_builder.remote_storage, S3Storage)
 
     # Note that this doesnt use pagination, so list is not guaranteed to be exhaustive.
+    assert neon_env_builder.remote_storage_client is not None
     response = neon_env_builder.remote_storage_client.list_objects_v2(
         Bucket=neon_env_builder.remote_storage.bucket_name,
         Prefix=prefix or neon_env_builder.remote_storage.prefix_in_bucket or "",
@@ -628,7 +629,7 @@ def test_timeline_delete_works_for_remote_smoke(
         )
 
     # for some reason the check above doesnt immediately take effect for the below.
-    # Assume it is mock server incosistency and check twice.
+    # Assume it is mock server inconsistency and check twice.
     wait_until(
         2,
         0.5,