diff --git a/flake.nix b/flake.nix
index 4cd1992f2..391ccbe98 100644
--- a/flake.nix
+++ b/flake.nix
@@ -62,6 +62,16 @@
 
         sfcgal = pkgs.callPackage ./nix/ext/sfcgal/sfcgal.nix { };
         pg_regress = pkgs.callPackage ./nix/ext/pg_regress.nix { };
+        pg_prove = pkgs.runCommand "pg_prove"
+          {
+            nativeBuildInputs = [ pkgs.makeWrapper ];
+          } ''
+          mkdir -p $out/bin
+          for x in pg_prove pg_tapgen; do
+            makeWrapper "${pkgs.perlPackages.TAPParserSourceHandlerpgTAP}/bin/$x" "$out/bin/$x" \
+              ${pkgs.lib.optionalString pkgs.stdenv.isLinux "--set LOCALE_ARCHIVE \"${pkgs.glibcLocales}/lib/locale/locale-archive\""}
+          done
+        '';
 
         # Our list of PostgreSQL extensions which come from upstream Nixpkgs.
         # These are maintained upstream and can easily be used here just by
@@ -323,25 +333,28 @@
             '';
 
           # Migrate between two data directories.
-          migrate-tool =
-            let
-              configFile = ./nix/tests/postgresql.conf.in;
-              getkeyScript = ./nix/tests/util/pgsodium_getkey.sh;
-              primingScript = ./nix/tests/prime.sql;
-              migrationData = ./nix/tests/migrations/data.sql;
-            in
-            pkgs.runCommand "migrate-postgres" { } ''
-              mkdir -p $out/bin
-              substitute ${./nix/tools/migrate-tool.sh.in} $out/bin/migrate-postgres \
-                --subst-var-by 'PSQL15_BINDIR' '${basePackages.psql_15.bin}' \
-                --subst-var-by 'PSQL_CONF_FILE' '${configFile}' \
-                --subst-var-by 'PGSODIUM_GETKEY' '${getkeyScript}' \
-                --subst-var-by 'PRIMING_SCRIPT' '${primingScript}' \
-                --subst-var-by 'MIGRATION_DATA' '${migrationData}'
-
-              chmod +x $out/bin/migrate-postgres
-            '';
-
+        migrate-tool =
+          let
+            configFile = ./nix/tests/postgresql.conf.in;
+            getkeyScript = ./nix/tests/util/pgsodium_getkey.sh;
+            primingScript = ./nix/tests/prime.sql;
+            migrationsDir = ./migrations;
+            pgupgradeTests = ./tests;
+            pgProve = pg_prove;
+          in
+          pkgs.runCommand "migrate-postgres" { } ''
+            mkdir -p $out/bin $out/migrations $out/tests
+            cp -r ${migrationsDir}/* $out/migrations
+            cp -r ${pgupgradeTests}/* $out/tests
+            substitute ${./nix/tools/migrate-tool.sh.in} $out/bin/migrate-postgres \
+              --subst-var-by 'PSQL_CONF_FILE' '${configFile}' \
+              --subst-var-by 'PGSODIUM_GETKEY' '${getkeyScript}' \
+              --subst-var-by 'PRIMING_SCRIPT' '${primingScript}' \
+              --subst-var-by 'MIGRATIONS_DIR' "$out/migrations"  \
+              --subst-var-by 'PGUPGRADE_TESTS' "$out/tests" \
+              --subst-var-by 'PG_PROVE' "${pgProve}"           
+            chmod +x $out/bin/migrate-postgres
+          '';
           start-replica = pkgs.runCommand "start-postgres-replica" { } ''
             mkdir -p $out/bin
             substitute ${./nix/tools/run-replica.sh.in} $out/bin/start-postgres-replica \
diff --git a/nix/ext/postgis.nix b/nix/ext/postgis.nix
index e0b6dfbeb..e76d9260e 100644
--- a/nix/ext/postgis.nix
+++ b/nix/ext/postgis.nix
@@ -61,6 +61,14 @@ stdenv.mkDerivation rec {
     ln -s ${postgresql}/bin/postgres $out/bin/postgres
   '';
 
+  installPhase = ''
+    mkdir -p $out/{bin,lib,share/postgresql/extension}
+    cp bin/*      $out/bin
+    cp *.so      $out/lib
+    cp *.sql     $out/share/postgresql/extension
+    cp *.control $out/share/postgresql/extension
+  '';
+  
   # create aliases for all commands adding version information
   postInstall = ''
     # Teardown the illusory postgres used for building; see postConfigure.
diff --git a/nix/tools/migrate-tool.sh.in b/nix/tools/migrate-tool.sh.in
index 94eef85a8..5a42503d9 100644
--- a/nix/tools/migrate-tool.sh.in
+++ b/nix/tools/migrate-tool.sh.in
@@ -1,123 +1,173 @@
 #!/usr/bin/env bash
 
-[ ! -z "$DEBUG" ] && set -x
-
-# first argument is the old version; a path 15 or 16
-if [[ $1 == /nix/store* ]]; then
-    if [ ! -L "$1/receipt.json" ] || [ ! -e "$1/receipt.json" ]; then
-        echo "ERROR: $1 does not look like a valid Postgres install"
+[ ! -z "$DEBUG" ] && set -eoux pipefail 
+
+# Function to build flake and return the output path
+build_flake() {
+    local flake_url="$1"
+    local temp_dir=$(mktemp -d)
+    if ! nix build "$flake_url" -o "$temp_dir/result"; then
+        echo "ERROR: Failed to build flake $flake_url"
         exit 1
     fi
-    OLDVER="$1"
-elif [ "$1" == "15" ]; then
-    PSQL15=@PSQL15_BINDIR@
-    OLDVER="$PSQL15"
-elif [ "$1" == "16" ]; then
-    PSQL16=@PSQL16_BINDIR@
-    OLDVER="$PSQL16"
-else
-    echo "Please provide a valid Postgres version (15 or 16), or a /nix/store path"
-    exit 1
-fi
+    echo "$temp_dir/result"
+}
 
-# second argument is the new version; 15 or 16
-if [[ $2 == /nix/store* ]]; then
-    if [ ! -L "$2/receipt.json" ] || [ ! -e "$2/receipt.json" ]; then
-        echo "ERROR: $1 does not look like a valid Postgres install"
-        exit 1
-    fi
-    NEWVER="$2"
-elif [ "$2" == "15" ]; then
-    PSQL15=@PSQL15_BINDIR@
-    NEWVER="$PSQL15"
-elif [ "$2" == "16" ]; then
-    PSQL16=@PSQL16_BINDIR@
-    NEWVER="$PSQL16"
-    echo "NEWVER IS $NEWVER"
-else
-    echo "Please provide a valid Postgres version (15 or 16), or a /nix/store path"
-    exit 1
-fi
+# First argument is the old version flake URL
+OLDVER=$(build_flake "$1")
+
+# Second argument is the new version flake URL
+NEWVER=$(build_flake "$2")
 
-# thid argument is the upgrade method: either pg_dumpall or pg_ugprade
+# Third argument is the upgrade method: either pg_dumpall or pg_upgrade
 if [ "$3" != "pg_dumpall" ] && [ "$3" != "pg_upgrade" ]; then
     echo "Please provide a valid upgrade method (pg_dumpall or pg_upgrade)"
     exit 1
 fi
 UPGRADE_METHOD="$3"
 
-echo "Old server build: PSQL $1"
-echo "New server build: PSQL $2"
+echo "Old server build: $OLDVER"
+echo "New server build: $NEWVER"
 echo "Upgrade method: $UPGRADE_METHOD"
 
-PORTNO="${2:-@PGSQL_DEFAULT_PORT@}"
+PORTNO="@PGSQL_DEFAULT_PORT@"
 DATDIR=$(mktemp -d)
 NEWDAT=$(mktemp -d)
+PGUSER=${PGUSER:-postgres}
+POSTGRES_PASSWORD=${POSTGRES_PASSWORD:-postgres}
+PG_PROVE=@PG_PROVE@
+PGUPGRADE_TESTS=@PGUPGRADE_TESTS@
 mkdir -p "$DATDIR" "$NEWDAT"
 
-echo "NOTE: using temporary directory $DATDIR for PSQL $1 data, which will not be removed"
-echo "NOTE: you are free to re-use this data directory at will"
+echo "using temporary directory $DATDIR for old data, which will not be removed"
+echo "you are free to re-use this data directory at will"
 echo
 
-$OLDVER/bin/initdb -D "$DATDIR" --locale=C
-$NEWVER/bin/initdb -D "$NEWDAT" --locale=C
+echo "PGUSER IS $PGUSER"
+
+$OLDVER/bin/initdb -U "$PGUSER" -D "$DATDIR" --locale=C
+$NEWVER/bin/initdb -U "$PGUSER" -D "$NEWDAT" --locale=C
 
 # NOTE (aseipp): we need to patch postgresql.conf to have the right pgsodium_getkey script
 PSQL_CONF_FILE=@PSQL_CONF_FILE@
 PGSODIUM_GETKEY_SCRIPT=@PGSODIUM_GETKEY@
-echo "NOTE: patching postgresql.conf files"
+echo "patching postgresql.conf files"
 for x in "$DATDIR" "$NEWDAT"; do
   sed \
     "s#@PGSODIUM_GETKEY_SCRIPT@#$PGSODIUM_GETKEY_SCRIPT#g" \
     $PSQL_CONF_FILE > "$x/postgresql.conf"
 done
 
-echo "NOTE: Starting first server (v${1}) to load data into the system"
+echo "Starting first server to load data into the system"
 $OLDVER/bin/pg_ctl start -D "$DATDIR"
 
 PRIMING_SCRIPT=@PRIMING_SCRIPT@
-MIGRATION_DATA=@MIGRATION_DATA@
+MIGRATIONS_DIR=@MIGRATIONS_DIR@
+
+echo "MIGRATIONS_DIR IS $MIGRATIONS_DIR"
+
+for sql in "$MIGRATIONS_DIR"/db/init-scripts/*.sql; do
+    echo "$0: running $sql"
+    $OLDVER/bin/psql -h localhost -d postgres -v ON_ERROR_STOP=1 --no-password --no-psqlrc -U postgres -f "$sql"
+done
 
-$OLDVER/bin/psql -h localhost -d postgres -Xf "$PRIMING_SCRIPT"
-$OLDVER/bin/psql -h localhost -d postgres -Xf "$MIGRATION_DATA"
+$OLDVER/bin/psql -h localhost -d postgres -v ON_ERROR_STOP=1 --no-password --no-psqlrc -U postgres -c "ALTER USER supabase_admin WITH PASSWORD '$PGPASSWORD'"
+# run migrations as super user - postgres user demoted in post-setup
+for sql in "$MIGRATIONS_DIR"/db/migrations/*.sql; do
+    echo "$0: running $sql"
+    $OLDVER/bin/psql -h localhost -d postgres -v ON_ERROR_STOP=1 --no-password --no-psqlrc -U supabase_admin -f "$sql"
+done
 
 if [ "$UPGRADE_METHOD" == "pg_upgrade" ]; then
-  echo "NOTE: Stopping old server (v${1}) to prepare for migration"
+  echo "Stopping old server"
   $OLDVER/bin/pg_ctl stop -D "$DATDIR"
 
-  echo "NOTE: Migrating old data $DATDIR to $NEWDAT using pg_upgrade"
+  echo "Starting old server"
+  $OLDVER/bin/pg_ctl start -D "$DATDIR" 
+
+  echo "Ensuring $PGUSER is a superuser in the old database"
+  "$OLDVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -c "ALTER USER $PGUSER WITH SUPERUSER;" || true
+
+  echo "Running pre-migration checks"
+  "$PG_PROVE"/bin/pg_prove --psql="$OLDVER"/bin/psql -h localhost -U supabase_admin -d postgres -p 5432 "$MIGRATIONS_DIR"/tests/test.sql
+
+  echo "Running fixtures"
+  "$OLDVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -f "$PGUPGRADE_TESTS/pg_upgrade/tests/97-enable-extensions.sql"
+  "$OLDVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -f "$PGUPGRADE_TESTS/pg_upgrade/tests/98-data-fixtures.sql"
+  "$OLDVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -f "$PGUPGRADE_TESTS/pg_upgrade/tests/99-fixtures.sql"
+
+  echo "Stopping old server"
+  $OLDVER/bin/pg_ctl stop -D "$DATDIR"
+
+  echo "Migrating old data $DATDIR to $NEWDAT using pg_upgrade"
 
   export PGDATAOLD="$DATDIR"
   export PGDATANEW="$NEWDAT"
   export PGBINOLD="$OLDVER/bin"
   export PGBINNEW="$NEWVER/bin"
 
-  if ! $NEWVER/bin/pg_upgrade --check; then
+  # Create a temporary directory for pg_upgrade to work in
+  UPGRADE_WORKDIR=$(mktemp -d)
+  echo "Using temporary directory for pg_upgrade: $UPGRADE_WORKDIR"
+
+  # Change to the temporary directory before running pg_upgrade
+  pushd "$UPGRADE_WORKDIR"
+
+  if ! $NEWVER/bin/pg_upgrade -U "$PGUSER" --check; then
       echo "ERROR: pg_upgrade check failed"
+      popd
       exit 1
   fi
 
-  echo "NOTE: pg_upgrade check passed, proceeding with migration"
-  $NEWVER/bin/pg_upgrade
-  rm -f delete_old_cluster.sh # we don't need this
+  echo "pg_upgrade check passed, proceeding with migration"
+  $NEWVER/bin/pg_upgrade -U "$PGUSER"
+
+  # Change back to the original directory
+  popd
+  echo "Migration complete, running post-migration checks"
+  echo "NEWDAT IS $NEWDAT"
+  $NEWVER/bin/pg_ctl start -D "$NEWDAT"
+  echo "Turning off JIT"
+  cat << EOF > "$NEWDAT"/jit_off.sql
+ALTER SYSTEM SET jit = off;
+SELECT pg_reload_conf();
+EOF
+  "$NEWVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -f "$NEWDAT"/jit_off.sql  
+
+  echo "Setting password encryption method to scram-sha-256"
+  "$NEWVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -c "ALTER SYSTEM SET password_encryption = 'scram-sha-256';"
+  "$NEWVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -c "SELECT pg_reload_conf();"
+
+  echo "Running post-migration data checks"
+  "$PG_PROVE"/bin/pg_prove --psql="$NEWVER"/bin/psql -h localhost -U supabase_admin -d postgres -p 5432 \
+    "$PGUPGRADE_TESTS/pg_upgrade/tests/01-schema.sql"
+  "$PG_PROVE"/bin/pg_prove --psql="$NEWVER"/bin/psql -h localhost -U supabase_admin -d postgres -p 5432 \
+    "$PGUPGRADE_TESTS/pg_upgrade/tests/02-data.sql"
+  "$PG_PROVE"/bin/pg_prove --psql="$NEWVER"/bin/psql -h localhost -U supabase_admin -d postgres -p 5432 \
+    "$PGUPGRADE_TESTS/pg_upgrade/tests/03-settings.sql"
+
+  # echo "Dumping database statistics for debugging"
+  # "$NEWVER"/bin/psql -h localhost -U supabase_admin -p 5432 -d postgres -c "SELECT schemaname, relname, n_live_tup FROM pg_stat_user_tables ORDER BY n_live_tup DESC;"
+
+  $NEWVER/bin/pg_ctl stop -D "$NEWDAT"
   exit 0
 fi
 
-if [ "$UPGRADE_METHOD" == "pg_dumpall" ]; then
-    SQLDAT="$DATDIR/dump.sql"
-    echo "NOTE: Exporting data via pg_dumpall ($SQLDAT)"
-    $NEWVER/bin/pg_dumpall -h localhost > "$SQLDAT"
+# if [ "$UPGRADE_METHOD" == "pg_dumpall" ]; then
+#     SQLDAT="$DATDIR/dump.sql"
+#     echo "Exporting data via pg_dumpall ($SQLDAT)"
+#     $NEWVER/bin/pg_dumpall -h localhost > "$SQLDAT"
 
-    echo "NOTE: Stopping old server (v${1}) to prepare for migration"
-    $OLDVER/bin/pg_ctl stop -D "$DATDIR"
+#     echo "Stopping old server (v${1}) to prepare for migration"
+#     $OLDVER/bin/pg_ctl stop -D "$DATDIR"
 
-    echo "NOTE: Starting second server (v${2}) to load data into the system"
-    $NEWVER/bin/pg_ctl start -D "$NEWDAT"
+#     echo "Starting second server (v${2}) to load data into the system"
+#     $NEWVER/bin/pg_ctl start -D "$NEWDAT"
 
-    echo "NOTE: Loading data into new server (v${2}) via 'cat | psql'"
-    cat "$SQLDAT" | $NEWVER/bin/psql -h localhost -d postgres
+#     echo "Loading data into new server (v${2}) via 'cat | psql'"
+#     cat "$SQLDAT" | $NEWVER/bin/psql -h localhost -d postgres
 
-    printf "\n\n\n\n"
-    echo "NOTE: Done, check logs. Stopping the server; new database is located at $NEWDAT"
-    $NEWVER/bin/pg_ctl stop -D "$NEWDAT"
-fi
+#     printf "\n\n\n\n"
+#     echo "Done, check logs. Stopping the server; new database is located at $NEWDAT"
+#     $NEWVER/bin/pg_ctl stop -D "$NEWDAT"
+# fi