feat: multiple versions for the wrappers extension

nix/checks.nix

@@ -307,12 +307,11 @@
             postgresql_17_src
             ;
         }
-        // pkgs.lib.optionalAttrs (system == "x86_64-linux") {
-          wrappers = import ./ext/tests/wrappers.nix {
+        // pkgs.lib.optionalAttrs (system == "x86_64-linux") (
+          import ./ext/tests {
             inherit self;
             inherit pkgs;
-          };
-          devShell = self'.devShells.default;
-        };
+          }
+        );
     };
 }

nix/ext/tests/default.nix

@@ -0,0 +1,173 @@
+{ self, pkgs }:
+let
+  testsDir = ./.;
+  testFiles = builtins.attrNames (builtins.readDir testsDir);
+  nixFiles = builtins.filter (
+    name: builtins.match ".*\\.nix$" name != null && name != "default.nix"
+  ) testFiles;
+  extTest =
+    extension_name:
+    let
+      pname = extension_name;
+      inherit (pkgs) lib;
+      installedExtension =
+        postgresMajorVersion: self.packages.${pkgs.system}."psql_${postgresMajorVersion}/exts/${pname}-all";
+      versions = postgresqlMajorVersion: (installedExtension postgresqlMajorVersion).versions;
+      postgresqlWithExtension =
+        postgresql:
+        let
+          majorVersion = lib.versions.major postgresql.version;
+          pkg = pkgs.buildEnv {
+            name = "postgresql-${majorVersion}-${pname}";
+            paths = [
+              postgresql
+              postgresql.lib
+              (installedExtension majorVersion)
+            ];
+            passthru = {
+              inherit (postgresql) version psqlSchema;
+              lib = pkg;
+              withPackages = _: pkg;
+            };
+            nativeBuildInputs = [ pkgs.makeWrapper ];
+            pathsToLink = [
+              "/"
+              "/bin"
+              "/lib"
+            ];
+            postBuild = ''
+              wrapProgram $out/bin/postgres --set NIX_PGLIBDIR $out/lib
+              wrapProgram $out/bin/pg_ctl --set NIX_PGLIBDIR $out/lib
+              wrapProgram $out/bin/pg_upgrade --set NIX_PGLIBDIR $out/lib
+            '';
+          };
+        in
+        pkg;
+    in
+    self.inputs.nixpkgs.lib.nixos.runTest {
+      name = pname;
+      hostPkgs = pkgs;
+      nodes.server =
+        { config, ... }:
+        {
+          virtualisation = {
+            forwardPorts = [
+              {
+                from = "host";
+                host.port = 13022;
+                guest.port = 22;
+              }
+            ];
+          };
+          services.openssh = {
+            enable = true;
+          };
+
+          services.postgresql = {
+            enable = true;
+            package = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_15;
+            enableTCPIP = true;
+            initialScript = pkgs.writeText "init-postgres-with-password" ''
+              CREATE USER test WITH PASSWORD 'secret';
+            '';
+            authentication = ''
+              host test postgres samenet scram-sha-256
+            '';
+          };
+
+          networking.firewall.allowedTCPPorts = [ config.services.postgresql.settings.port ];
+
+          specialisation.postgresql17.configuration = {
+            services.postgresql = {
+              package = lib.mkForce (postgresqlWithExtension self.packages.${pkgs.system}.postgresql_17);
+            };
+
+            systemd.services.postgresql-migrate = {
+              serviceConfig = {
+                Type = "oneshot";
+                RemainAfterExit = true;
+                User = "postgres";
+                Group = "postgres";
+                StateDirectory = "postgresql";
+                WorkingDirectory = "${builtins.dirOf config.services.postgresql.dataDir}";
+              };
+              script =
+                let
+                  oldPostgresql = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_15;
+                  newPostgresql = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_17;
+                  oldDataDir = "${builtins.dirOf config.services.postgresql.dataDir}/${oldPostgresql.psqlSchema}";
+                  newDataDir = "${builtins.dirOf config.services.postgresql.dataDir}/${newPostgresql.psqlSchema}";
+                in
+                ''
+                  if [[ ! -d ${newDataDir} ]]; then
+                    install -d -m 0700 -o postgres -g postgres "${newDataDir}"
+                    ${newPostgresql}/bin/initdb -D "${newDataDir}"
+                    ${newPostgresql}/bin/pg_upgrade --old-datadir "${oldDataDir}" --new-datadir "${newDataDir}" \
+                      --old-bindir "${oldPostgresql}/bin" --new-bindir "${newPostgresql}/bin"
+                  else
+                    echo "${newDataDir} already exists"
+                  fi
+                '';
+            };
+
+            systemd.services.postgresql = {
+              after = [ "postgresql-migrate.service" ];
+              requires = [ "postgresql-migrate.service" ];
+            };
+          };
+        };
+      testScript =
+        { nodes, ... }:
+        let
+          pg17-configuration = "${nodes.server.system.build.toplevel}/specialisation/postgresql17";
+        in
+        ''
+          versions = {
+            "15": [${lib.concatStringsSep ", " (map (s: ''"${s}"'') (versions "15"))}],
+            "17": [${lib.concatStringsSep ", " (map (s: ''"${s}"'') (versions "17"))}],
+          }
+          extension_name = "${pname}"
+          support_upgrade = True
+          pg17_configuration = "${pg17-configuration}"
+
+          ${builtins.readFile ./lib.py}
+
+          start_all()
+
+          server.wait_for_unit("multi-user.target")
+          server.wait_for_unit("postgresql.service")
+
+          test = PostgresExtensionTest(server, extension_name, versions, support_upgrade)
+
+          with subtest("Check upgrade path with postgresql 15"):
+            test.check_upgrade_path("15")
+
+          last_version = None
+          with subtest("Check the install of the last version of the extension"):
+            last_version = test.check_install_last_version("15")
+
+          with subtest("switch to postgresql 17"):
+            server.succeed(
+              f"{pg17_configuration}/bin/switch-to-configuration test >&2"
+            )
+
+          with subtest("Check last version of the extension after upgrade"):
+            test.assert_version_matches(last_version)
+
+          with subtest("Check upgrade path with postgresql 17"):
+            test.check_upgrade_path("17")
+        '';
+    };
+in
+builtins.listToAttrs (
+  map (file: {
+    name = "ext-" + builtins.replaceStrings [ ".nix" ] [ "" ] file;
+    value = import (testsDir + "/${file}") { inherit self pkgs; };
+  }) nixFiles
+)
+// builtins.listToAttrs (
+  map (extName: {
+    name = "ext-${extName}";
+    value = extTest extName;
+  }) [ "wrappers" ]
+)

nix/ext/tests/lib.py

@@ -0,0 +1,109 @@
+"""PostgreSQL extension testing framework for multi-version compatibility.
+
+This module provides a test framework for PostgreSQL extensions that need to be
+tested across multiple PostgreSQL versions and extension versions. It handles
+installation, upgrades, and version verification of PostgreSQL extensions.
+"""
+
+from typing import Sequence, Mapping
+from test_driver.machine import Machine
+
+Versions = Mapping[str, Sequence[str]]
+
+class PostgresExtensionTest(object):
+    def __init__(self, vm: Machine, extension_name: str, versions: Versions, support_upgrade: bool = True):
+        """Initialize the PostgreSQL extension test framework.
+
+        Args:
+            vm: Test machine instance for executing commands
+            extension_name: Name of the PostgreSQL extension to test
+            versions: Mapping of PostgreSQL versions to available extension versions
+            support_upgrade: Whether the extension supports in-place upgrades
+        """
+        self.vm = vm
+        self.extension_name = extension_name
+        self.versions = versions
+        self.support_upgrade = support_upgrade
+
+    def run_sql(self, query: str) -> str:
+        return self.vm.succeed(f"""sudo -u postgres psql -t -A -F\",\" -c \"{query}\" """).strip()
+
+    def drop_extension(self):
+        self.run_sql(f"DROP EXTENSION IF EXISTS {self.extension_name};")
+
+    def install_extension(self, version: str):
+        self.run_sql(f"""CREATE EXTENSION {self.extension_name} WITH VERSION '{version}' CASCADE;""")
+        # Verify version was installed correctly
+        self.assert_version_matches(version)
+
+    def update_extension(self, version: str):
+        self.run_sql(f"""ALTER EXTENSION {self.extension_name} UPDATE TO '{version}';""")
+        # Verify version was installed correctly
+        self.assert_version_matches(version)
+
+    def get_installed_version(self) -> str:
+        """Get the currently installed version of the extension.
+
+        Returns:
+            Version string of the currently installed extension,
+            or empty string if extension is not installed
+        """
+        return self.run_sql(f"""SELECT extversion FROM pg_extension WHERE extname = '{self.extension_name}';""")
+
+    def assert_version_matches(self, expected_version: str):
+        """Check if the installed version matches the expected version.
+
+        Args:
+            expected_version: Expected version string to verify against
+
+        Raises:
+            AssertionError: If the installed version does not match the expected version
+        """
+        installed_version = self.get_installed_version()
+        assert installed_version == expected_version, f"Expected version {expected_version}, but found {installed_version}"
+
+    def check_upgrade_path(self, pg_version):
+        """Test the complete upgrade path for a PostgreSQL version.
+
+        This method tests all available extension versions for a given PostgreSQL
+        version, either through in-place upgrades or reinstallation depending on
+        the support_upgrade setting.
+
+        Args:
+            pg_version: PostgreSQL version to test (e.g., "14", "15")
+
+        Raises:
+            ValueError: If no versions are available for the specified PostgreSQL version
+            AssertionError: If version installation or upgrade fails
+        """
+        available_versions = self.versions.get(pg_version, [])
+        if not available_versions:
+            raise ValueError(f"No versions available for PostgreSQL version {pg_version}")
+
+        # Install and verify first version
+        firstVersion = available_versions[0]
+        self.drop_extension()
+        self.install_extension(firstVersion)
+
+        # Test remaining versions
+        for version in available_versions[1:]:
+            if self.support_upgrade:
+                self.update_extension(version)
+            else:
+                self.drop_extension()
+                self.install_extension(version)
+
+
+    def check_install_last_version(self, pg_version: str) -> str:
+        """Test if the install of the last version of the extension works for a given PostgreSQL version.
+
+        Args:
+            pg_version: PostgreSQL version to check (e.g., "14", "15")
+        """
+        available_versions = self.versions.get(pg_version, [])
+        if not available_versions:
+            raise ValueError(f"No versions available for PostgreSQL version {pg_version}")
+        last_version = available_versions[-1]
+        self.drop_extension()
+        self.install_extension(last_version)
+        return last_version