feat: switch both timescaledb and its loader when switching versions

Also reuse default tests defined in lib.py to test extension upgrade
paths and switching extension versions.

Author: jfroche

nix/ext/tests/lib.py

@@ -147,7 +147,7 @@ def check_switch_extension_with_background_worker(
             )
         last_version = available_versions[-1]
         assert ext_version.endswith(
-            f"{self.extension_name}-{last_version}.so"
+            f"{last_version}.so"
         ), f"Expected {self.extension_name} version {last_version}, but found {ext_version}"
 
         # Switch to the first version
@@ -157,13 +157,13 @@ def check_switch_extension_with_background_worker(
         # Check that we are using the first version now
         ext_version = self.vm.succeed(f"readlink -f {extension_lib_path}").strip()
         assert ext_version.endswith(
-            f"{self.extension_name}-{first_version}.so"
+            f"{first_version}.so"
         ), f"Expected {self.extension_name} version {first_version}, but found {ext_version}"
 
         # Switch to the first version
         self.vm.succeed(f"switch_{self.extension_name}_version {last_version}")
         # Check that we are using the last version now
         ext_version = self.vm.succeed(f"readlink -f {extension_lib_path}").strip()
         assert ext_version.endswith(
-            f"{self.extension_name}-{last_version}.so"
+            f"{last_version}.so"
         ), f"Expected {self.extension_name} version {last_version}, but found {ext_version}"

nix/ext/tests/timescaledb.nix

@@ -1,17 +1,16 @@
 { self, pkgs }:
 let
+  pname = "timescaledb";
   inherit (pkgs) lib;
   installedExtension =
-    postgresMajorVersion:
-    self.packages.${pkgs.system}."psql_${postgresMajorVersion}/exts/timescaledb-all";
+    postgresMajorVersion: self.packages.${pkgs.system}."psql_${postgresMajorVersion}/exts/${pname}-all";
   versions = (installedExtension "15").versions;
-  firstVersion = lib.head versions;
   postgresqlWithExtension =
     postgresql:
     let
       majorVersion = lib.versions.major postgresql.version;
       pkg = pkgs.buildEnv {
-        name = "postgresql-${majorVersion}-timescaledb";
+        name = "postgresql-${majorVersion}-${pname}";
         paths = [
           postgresql
           postgresql.lib
@@ -36,65 +35,44 @@ let
       };
     in
     pkg;
+  psql_15 = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_15;
 in
 self.inputs.nixpkgs.lib.nixos.runTest {
   name = "timescaledb";
   hostPkgs = pkgs;
   nodes.server =
     { ... }:
     {
-      virtualisation = {
-        forwardPorts = [
-          {
-            from = "host";
-            host.port = 13022;
-            guest.port = 22;
-          }
-        ];
-      };
-      services.openssh = {
-        enable = true;
-      };
-      users.users.root.openssh.authorizedKeys.keys = [
-        "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIArkmq6Th79Z4klW6Urgi4phN8yq769/l/10jlE00tU9"
-      ];
-
       services.postgresql = {
         enable = true;
-        package = postgresqlWithExtension self.packages.${pkgs.system}.postgresql_15;
+        package = (postgresqlWithExtension psql_15);
         settings = {
           shared_preload_libraries = "timescaledb";
         };
       };
-
-      specialisation.postgresql15.configuration = {
-        services.postgresql = {
-          package = lib.mkForce (postgresqlWithExtension self.packages.${pkgs.system}.postgresql_15);
-        };
-      };
     };
   testScript =
     { ... }:
     ''
-      def run_sql(query):
-          return server.succeed(f"""sudo -u postgres psql -t -A -F\",\" -c \"{query}\" """).strip()
-
-      def check_upgrade_path():
-          with subtest("Check timescaledb upgrade path"):
-              server.succeed("sudo -u postgres psql -c 'DROP EXTENSION IF EXISTS timescaledb;'")
-              run_sql(r"""CREATE EXTENSION timescaledb WITH VERSION \"${firstVersion}\";""")
-              installed_version = run_sql(r"""SELECT extversion FROM pg_extension WHERE extname = 'timescaledb';""")
-              assert installed_version == "${firstVersion}", f"Expected timescaledb version ${firstVersion}, but found {installed_version}"
-              for version in [${lib.concatStringsSep ", " (map (s: ''"${s}"'') versions)}][1:]:
-                  run_sql(f"""ALTER EXTENSION timescaledb UPDATE TO '{version}';""")
-                  installed_version = run_sql(r"""SELECT extversion FROM pg_extension WHERE extname = 'timescaledb';""")
-                  assert installed_version == version, f"Expected timescaledb version {version}, but found {installed_version}"
+      ${builtins.readFile ./lib.py}
 
       start_all()
 
       server.wait_for_unit("multi-user.target")
       server.wait_for_unit("postgresql.service")
 
-      check_upgrade_path()
+      versions = {
+        "15": [${lib.concatStringsSep ", " (map (s: ''"${s}"'') versions)}],
+      }
+      extension_name = "${pname}"
+      support_upgrade = True
+
+      test = PostgresExtensionTest(server, extension_name, versions, support_upgrade)
+
+      with subtest("Check upgrade path with postgresql 15"):
+        test.check_upgrade_path("15")
+
+      with subtest("Test switch_${pname}_version"):
+        test.check_switch_extension_with_background_worker(Path("${psql_15}/lib/${pname}.so"), "15")
     '';
 }

nix/ext/timescaledb.nix

@@ -9,6 +9,8 @@
   buildEnv,
   makeWrapper,
   switch-ext-version,
+  coreutils,
+  writeShellApplication,
 }:
 
 let
@@ -65,7 +67,7 @@ let
 
         # Rename the loader library to be version-specific
         if [ -f $out/lib/timescaledb${postgresql.dlSuffix} ]; then
-          mv $out/lib/timescaledb${postgresql.dlSuffix} $out/lib/timescaledb-${version}-loader${postgresql.dlSuffix}
+          mv $out/lib/timescaledb${postgresql.dlSuffix} $out/lib/timescaledb-loader-${version}${postgresql.dlSuffix}
         fi
 
         # The versioned library (timescaledb-VERSION.so) is already correctly named
@@ -98,6 +100,16 @@ let
   packages = builtins.attrValues (
     lib.mapAttrs (name: value: build name value.hash (value.revision or name)) supportedVersions
   );
+  switch-timescaledb-loader = writeShellApplication {
+    name = "switch_timescaledb_loader";
+    runtimeInputs = [ coreutils ];
+    text = ''
+      EXT_LOADER_TO_USE="$EXT_WRAPPER_LIB/$EXT_NAME-loader-$VERSION${postgresql.dlSuffix}"
+      if [ -f "$EXT_LOADER_TO_USE" ]; then
+        ln -sfnv "$EXT_LOADER_TO_USE" "$EXT_WRAPPER_LIB/$EXT_NAME-loader${postgresql.dlSuffix}"
+      fi
+    '';
+  };
 in
 buildEnv {
   name = pname;
@@ -110,7 +122,7 @@ buildEnv {
     } > $out/share/postgresql/extension/${pname}.control
 
     # Create symlink from the latest versioned loader to timescaledb.so
-    ln -sfn ${pname}-${latestVersion}-loader${postgresql.dlSuffix} $out/lib/${pname}${postgresql.dlSuffix}
+    ln -sfn ${pname}-loader-${latestVersion}${postgresql.dlSuffix} $out/lib/${pname}${postgresql.dlSuffix}
 
     # The versioned extension libraries (timescaledb-VERSION.so) are already in place
 
@@ -119,13 +131,14 @@ buildEnv {
        test "$(ls -A $out/lib/${pname}*${postgresql.dlSuffix} | wc -l)" -gt 0
     )
     makeWrapper ${lib.getExe switch-ext-version} $out/bin/switch_timescaledb_version \
-      --prefix EXT_WRAPPER : "$out" --prefix EXT_NAME : "${pname}"
-
+      --prefix EXT_WRAPPER : "$out" --prefix EXT_NAME : "${pname}" --prefix EXTRA_STEPS : ${lib.getExe switch-timescaledb-loader}
   '';
+
   pathsToLink = [
     "/lib"
     "/share/postgresql/extension"
   ];
+
   passthru = {
     inherit versions numberOfVersions switch-ext-version;
     pname = "${pname}-all";

nix/packages/switch-ext-version.nix

@@ -54,6 +54,12 @@ writeShellApplication {
     # Update library symlink
     ln -sfnv "$EXT_LIB_TO_USE" "$EXT_WRAPPER_LIB/$EXT_NAME${postgresql.dlSuffix}"
 
+    # Handle extension specific steps
+    if [ -x "''${EXTRA_STEPS:-}" ]; then
+      #shellcheck disable=SC1090
+      source "''${EXTRA_STEPS}"
+    fi
+
     # Update control file
     EXT_WRAPPER_SHARE="$EXT_WRAPPER/share/postgresql/extension"
     echo "default_version = '$VERSION'" > "$EXT_WRAPPER_SHARE/$EXT_NAME.control"