Fetch remote extension in ALTER EXTENSION UPDATE statements (#11102)

Previously, remote extensions were not fetched unless they were used in
some other manner. For instance, loading a BM25 index in pg_search
fetches the pg_search extension. However, if on a fresh compute with
pg_search 0.15.5 installed, the user ran `ALTER EXTENSION pg_search
UPDATE TO '0.15.6'` without first using the pg_search extension, we
would not fetch the extension and fail to find an update path.

Signed-off-by: Tristan Partin <tristan@neon.tech>

compute_tools/src/extension_server.rs

@@ -202,8 +202,24 @@ pub async fn download_extension(
     // move contents of the libdir / sharedir in unzipped archive to the correct local paths
     for paths in [sharedir_paths, libdir_paths] {
         let (zip_dir, real_dir) = paths;
+
+        let dir = match std::fs::read_dir(&zip_dir) {
+            Ok(dir) => dir,
+            Err(e) => match e.kind() {
+                // In the event of a SQL-only extension, there would be nothing
+                // to move from the lib/ directory, so note that in the log and
+                // move on.
+                std::io::ErrorKind::NotFound => {
+                    info!("nothing to move from {}", zip_dir);
+                    continue;
+                }
+                _ => return Err(anyhow::anyhow!(e)),
+            },
+        };
+
         info!("mv {zip_dir:?}/*  {real_dir:?}");
-        for file in std::fs::read_dir(zip_dir)? {
+
+        for file in dir {
             let old_file = file?.path();
             let new_file =
                 Path::new(&real_dir).join(old_file.file_name().context("error parsing file")?);

test_runner/regress/data/extension_test/5670669815/v14/ext_index.json

@@ -1,24 +0,0 @@
-{
-    "public_extensions": [
-        "anon",
-        "pg_buffercache"
-    ],
-    "library_index": {
-        "anon": "anon",
-        "pg_buffercache": "pg_buffercache"
-    },
-    "extension_data": {
-        "pg_buffercache": {
-            "control_data": {
-                "pg_buffercache.control": "# pg_buffercache extension \ncomment = 'examine the shared buffer cache' \ndefault_version = '1.3' \nmodule_pathname = '$libdir/pg_buffercache' \nrelocatable = true \ntrusted=true"
-            },
-            "archive_path": "5670669815/v14/extensions/pg_buffercache.tar.zst"
-        },
-        "anon": {
-            "control_data": {
-                "anon.control": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n"
-            },
-            "archive_path": "5670669815/v14/extensions/anon.tar.zst"
-        }
-    }
-}

test_runner/regress/data/extension_test/5670669815/v15/ext_index.json

@@ -1,17 +0,0 @@
-{
-    "public_extensions": [
-        "anon"
-    ],
-    "library_index": {
-        "anon": "anon"
-    },
-    "extension_data": {
-        "anon": {
-            "control_data": {
-                "anon.control": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n"
-            },
-            "archive_path": "5670669815/v15/extensions/anon.tar.zst"
-        }
-    }
-}
-

test_runner/regress/data/extension_test/5670669815/v16/ext_index.json

@@ -1,7 +0,0 @@
-{
-    "public_extensions": [],
-    "library_index": {
-        "TODO": "We still need PG16 extensions"
-    },
-    "extension_data": {}
-}

test_runner/regress/data/extension_test/5670669815/v17/ext_index.json

@@ -1,7 +0,0 @@
-{
-    "public_extensions": [],
-    "library_index": {
-        "TODO": "We still need PG17 extensions"
-    },
-    "extension_data": {}
-}

test_runner/regress/data/test_remote_extensions/test_extension/sql/test_extension--1.0--1.1.sql

@@ -0,0 +1,10 @@
+\echo Use "ALTER EXTENSION test_extension UPDATE TO '1.1'" to load this file. \quit
+
+CREATE FUNCTION test_extension.fun_fact()
+RETURNS void
+IMMUTABLE LEAKPROOF PARALLEL SAFE
+AS $$
+BEGIN
+    RAISE NOTICE 'Neon has a melting point of -246.08 C';
+END;
+$$ LANGUAGE 'plpgsql';

test_runner/regress/data/test_remote_extensions/test_extension/sql/test_extension--1.0.sql

@@ -0,0 +1,12 @@
+\echo Use "CREATE EXTENSION test_extension" to load this file. \quit
+
+CREATE SCHEMA test_extension;
+
+CREATE FUNCTION test_extension.motd()
+RETURNS void
+IMMUTABLE LEAKPROOF PARALLEL SAFE
+AS $$
+BEGIN
+    RAISE NOTICE 'Have a great day';
+END;
+$$ LANGUAGE 'plpgsql';

test_runner/regress/data/test_remote_extensions/test_extension/test_extension.control

@@ -0,0 +1 @@
+comment = 'Test extension'

test_runner/regress/test_download_extensions.py

@@ -2,26 +2,26 @@ from __future__ import annotations
 
 import os
 import shutil
-from contextlib import closing
+import tarfile
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import pytest
+import zstandard
 from fixtures.log_helper import log
 from fixtures.metrics import parse_metrics
-from fixtures.neon_fixtures import (
-    NeonEnvBuilder,
-)
-from fixtures.pg_version import PgVersion
-from fixtures.utils import skip_on_postgres
 from pytest_httpserver import HTTPServer
-from werkzeug.wrappers.request import Request
 from werkzeug.wrappers.response import Response
 
 if TYPE_CHECKING:
     from typing import Any
 
     from fixtures.httpserver import ListenAddress
+    from fixtures.neon_fixtures import (
+        NeonEnvBuilder,
+    )
+    from fixtures.pg_version import PgVersion
+    from werkzeug.wrappers.request import Request
 
 
 # use neon_env_builder_local fixture to override the default neon_env_builder fixture
@@ -31,13 +31,13 @@ def neon_env_builder_local(
     neon_env_builder: NeonEnvBuilder,
     test_output_dir: Path,
     pg_distrib_dir: Path,
-    pg_version: PgVersion,
 ) -> NeonEnvBuilder:
     test_local_pginstall = test_output_dir / "pg_install"
     log.info(f"copy {pg_distrib_dir} to {test_local_pginstall}")
-    shutil.copytree(
-        pg_distrib_dir / pg_version.v_prefixed, test_local_pginstall / pg_version.v_prefixed
-    )
+
+    # We can't copy only the version that we are currently testing because other
+    # binaries like the storage controller need specific Postgres versions.
+    shutil.copytree(pg_distrib_dir, test_local_pginstall)
 
     neon_env_builder.pg_distrib_dir = test_local_pginstall
     log.info(f"local neon_env_builder.pg_distrib_dir: {neon_env_builder.pg_distrib_dir}")
@@ -45,89 +45,92 @@ def neon_env_builder_local(
     return neon_env_builder
 
 
-@skip_on_postgres(PgVersion.V16, reason="TODO: PG16 extension building")
-@skip_on_postgres(PgVersion.V17, reason="TODO: PG17 extension building")
 def test_remote_extensions(
     httpserver: HTTPServer,
     neon_env_builder_local: NeonEnvBuilder,
     httpserver_listen_address: ListenAddress,
+    test_output_dir: Path,
+    base_dir: Path,
     pg_version: PgVersion,
 ):
-    # setup mock http server
-    # that expects request for anon.tar.zst
-    # and returns the requested file
+    # Setup a mock nginx S3 gateway which will return our test extension.
     (host, port) = httpserver_listen_address
     extensions_endpoint = f"http://{host}:{port}/pg-ext-s3-gateway"
 
     build_tag = os.environ.get("BUILD_TAG", "latest")
-    archive_path = f"{build_tag}/v{pg_version}/extensions/anon.tar.zst"
+    archive_route = f"{build_tag}/v{pg_version}/extensions/test_extension.tar.zst"
+    tarball = test_output_dir / "test_extension.tar"
+    extension_dir = (
+        base_dir / "test_runner" / "regress" / "data" / "test_remote_extensions" / "test_extension"
+    )
 
-    def endpoint_handler_build_tag(request: Request) -> Response:
+    # Create tarball
+    with tarfile.open(tarball, "x") as tarf:
+        tarf.add(
+            extension_dir / "sql" / "test_extension--1.0.sql",
+            arcname="share/extension/test_extension--1.0.sql",
+        )
+        tarf.add(
+            extension_dir / "sql" / "test_extension--1.0--1.1.sql",
+            arcname="share/extension/test_extension--1.0--1.1.sql",
+        )
+
+    def handler(request: Request) -> Response:
         log.info(f"request: {request}")
 
-        file_name = "anon.tar.zst"
-        file_path = f"test_runner/regress/data/extension_test/5670669815/v{pg_version}/extensions/anon.tar.zst"
-        file_size = os.path.getsize(file_path)
-        fh = open(file_path, "rb")
+        # Compress tarball
+        compressor = zstandard.ZstdCompressor()
+        with open(tarball, "rb") as f:
+            compressed_data = compressor.compress(f.read())
 
         return Response(
-            fh,
+            compressed_data,
             mimetype="application/octet-stream",
             headers=[
-                ("Content-Length", str(file_size)),
-                ("Content-Disposition", f'attachment; filename="{file_name}"'),
+                ("Content-Length", str(len(compressed_data))),
             ],
             direct_passthrough=True,
         )
 
     httpserver.expect_request(
-        f"/pg-ext-s3-gateway/{archive_path}", method="GET"
-    ).respond_with_handler(endpoint_handler_build_tag)
+        f"/pg-ext-s3-gateway/{archive_route}", method="GET"
+    ).respond_with_handler(handler)
 
     # Start a compute node with remote_extension spec
     # and check that it can download the extensions and use them to CREATE EXTENSION.
     env = neon_env_builder_local.init_start()
     env.create_branch("test_remote_extensions")
-    endpoint = env.endpoints.create(
-        "test_remote_extensions",
-        config_lines=["log_min_messages=debug3"],
-    )
+    endpoint = env.endpoints.create("test_remote_extensions")
+
+    with open(extension_dir / "test_extension.control", encoding="utf-8") as f:
+        control_data = f.read()
 
     # mock remote_extensions spec
     spec: dict[str, Any] = {
-        "public_extensions": ["anon"],
+        "public_extensions": ["test_extension"],
         "custom_extensions": None,
         "library_index": {
-            "anon": "anon",
+            "test_extension": "test_extension",
         },
         "extension_data": {
-            "anon": {
+            "test_extension": {
                 "archive_path": "",
                 "control_data": {
-                    "anon.control": "# PostgreSQL Anonymizer (anon) extension\ncomment = 'Data anonymization tools'\ndefault_version = '1.1.0'\ndirectory='extension/anon'\nrelocatable = false\nrequires = 'pgcrypto'\nsuperuser = false\nmodule_pathname = '$libdir/anon'\ntrusted = true\n"
+                    "test_extension.control": control_data,
                 },
             },
         },
     }
-    spec["extension_data"]["anon"]["archive_path"] = archive_path
 
     endpoint.create_remote_extension_spec(spec)
 
-    endpoint.start(
-        remote_ext_config=extensions_endpoint,
-    )
+    endpoint.start(remote_ext_config=extensions_endpoint)
 
-    # this is expected to fail if there's no pgcrypto extension, that's ok
-    # we just want to check that the extension was downloaded
-    try:
-        with closing(endpoint.connect()) as conn:
-            with conn.cursor() as cur:
-                # Check that appropriate files were downloaded
-                cur.execute("CREATE EXTENSION anon")
-                res = [x[0] for x in cur.fetchall()]
-                log.info(res)
-    except Exception as err:
-        assert "pgcrypto" in str(err), f"unexpected error creating anon extension {err}"
+    with endpoint.connect() as conn:
+        with conn.cursor() as cur:
+            # Check that appropriate files were downloaded
+            cur.execute("CREATE EXTENSION test_extension VERSION '1.0'")
+            cur.execute("SELECT test_extension.motd()")
 
     httpserver.check()
 
@@ -138,7 +141,47 @@ def test_remote_extensions(
     remote_ext_requests = metrics.query_all(
         "compute_ctl_remote_ext_requests_total",
         # Check that we properly report the filename in the metrics
-        {"filename": "anon.tar.zst"},
+        {"filename": "test_extension.tar.zst"},
+    )
+    assert len(remote_ext_requests) == 1
+    for sample in remote_ext_requests:
+        assert sample.value == 1
+
+    endpoint.stop()
+
+    # Remove the extension files to force a redownload of the extension.
+    for file in (
+        "test_extension.control",
+        "test_extension--1.0.sql",
+        "test_extension--1.0--1.1.sql",
+    ):
+        (
+            test_output_dir
+            / "pg_install"
+            / f"v{pg_version}"
+            / "share"
+            / "postgresql"
+            / "extension"
+            / file
+        ).unlink()
+
+    endpoint.start(remote_ext_config=extensions_endpoint)
+
+    # Test that ALTER EXTENSION UPDATE statements also fetch remote extensions.
+    with endpoint.connect() as conn:
+        with conn.cursor() as cur:
+            # Check that appropriate files were downloaded
+            cur.execute("ALTER EXTENSION test_extension UPDATE TO '1.1'")
+            cur.execute("SELECT test_extension.fun_fact()")
+
+    # Check that we properly recorded downloads in the metrics
+    client = endpoint.http_client()
+    raw_metrics = client.metrics()
+    metrics = parse_metrics(raw_metrics)
+    remote_ext_requests = metrics.query_all(
+        "compute_ctl_remote_ext_requests_total",
+        # Check that we properly report the filename in the metrics
+        {"filename": "test_extension.tar.zst"},
     )
     assert len(remote_ext_requests) == 1
     for sample in remote_ext_requests:

vendor/revisions.json

@@ -1,18 +1,18 @@
 {
   "v17": [
     "17.4",
-    "780efda2ef8d629495cc289624534ba8cde40779"
+    "e5e87b9f52d0eaeb83f3e2517bb9727aac37729b"
   ],
   "v16": [
     "16.8",
-    "78050f965f2e550fd6e58f837394cb3d080d7d42"
+    "512856aaa8bedbaa8f06811449518dcb0c2e5d8f"
   ],
   "v15": [
     "15.12",
-    "533be42f7da97e614ce1c494fafe3e49f53991b1"
+    "ee794ba767eef9b10260ef67d3a58084f1dabd6f"
   ],
   "v14": [
     "14.17",
-    "b1425505c6f9a622a5aadf3ee362740519993310"
+    "7b7592e74059f795b64f06860cea97673418f35e"
   ]
 }