diff --git a/modules/nixos/services/postgresql/default.nix b/modules/nixos/services/postgresql/default.nix
index 6f51f3e..1c06736 100644
--- a/modules/nixos/services/postgresql/default.nix
+++ b/modules/nixos/services/postgresql/default.nix
@@ -27,17 +27,26 @@ in
 
       environment.systemPackages =
         let
-          newpg = config.containers.temp-pg.config.services.postgresql;
+          pgCfg = config.services.postgresql;
+          newPackage' = pkgs.postgresql_13;
+
+          oldPackage = if pgCfg.enableJIT then pgCfg.package.withJIT else pgCfg.package;
+          oldData = pgCfg.dataDir;
+          oldBin = "${if pgCfg.extraPlugins == [] then oldPackage else oldPackage.withPackages pgCfg.extraPlugins}/bin";
+
+          newPackage = if pgCfg.enableJIT then newPackage'.withJIT else newPackage';
+          newData = "/var/lib/postgresql/${newPackage.psqlSchema}";
+          newBin = "${if pgCfg.extraPlugins == [] then newPackage else newPackage.withPackages pgCfg.extraPlugins}/bin";
         in
         [
           (pkgs.writeScriptBin "upgrade-pg-cluster" ''
             #!/usr/bin/env bash
 
-            set -x
-            export OLDDATA="${config.services.postgresql.dataDir}"
-            export NEWDATA="${newpg.dataDir}"
-            export OLDBIN="${config.services.postgresql.package}/bin"
-            export NEWBIN="${newpg.package}/bin"
+            set -eux
+            export OLDDATA="${oldData}"
+            export NEWDATA="${newData}"
+            export OLDBIN="${oldBin}"
+            export NEWBIN="${newBin}"
 
             if [ "$OLDDATA" -ef "$NEWDATA" ]; then
               echo "Cannot migrate to same data directory" >&2
@@ -46,13 +55,13 @@ in
 
             install -d -m 0700 -o postgres -g postgres "$NEWDATA"
             cd "$NEWDATA"
-            sudo -u postgres $NEWBIN/initdb -D "$NEWDATA"
+            sudo -u postgres "$NEWBIN/initdb" -D "$NEWDATA"
 
             systemctl stop postgresql    # old one
 
-            sudo -u postgres $NEWBIN/pg_upgrade \
+            sudo -u postgres "$NEWBIN/pg_upgrade" \
               --old-datadir "$OLDDATA" --new-datadir "$NEWDATA" \
-              --old-bindir $OLDBIN --new-bindir $NEWBIN \
+              --old-bindir "$OLDBIN" --new-bindir "$NEWBIN" \
               "$@"
           '')
         ];