From 2f349ce456d1fadcee0c6ead9c68d0bc00d0dcec Mon Sep 17 00:00:00 2001
From: eaydingol <ebruaydin@gmail.com>
Date: Fri, 5 Dec 2025 14:18:06 +0300
Subject: [PATCH] Support loading citus and columnar so files from a given
 directory.

---
 src/test/regress/Makefile            |  6 ++-
 src/test/regress/pg_regress_multi.pl | 56 +++++++++++++++++++++++++++-
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/src/test/regress/Makefile b/src/test/regress/Makefile
index ba5cf58b1..9b9c0425a 100644
--- a/src/test/regress/Makefile
+++ b/src/test/regress/Makefile
@@ -13,6 +13,10 @@ endif
 ifndef CITUSVERSION
 CITUSVERSION := ""
 endif
+
+ifndef CITUSLIBDIR
+CITUSLIBDIR := ""
+endif
 # Add ./bin to $PATH so we use our custom diff instead of the default diff tool.
 # We do this to be able to mask shard Ids, placement Ids, node ports, etc.
 MAKEFILE_DIR := $(dir $(realpath $(firstword $(MAKEFILE_LIST))))
@@ -31,7 +35,7 @@ export PGISOLATIONTIMEOUT = 60
 ## Citus regression support
 ##
 MULTI_INSTALLDIR=$(CURDIR)/tmp_check/install
-pg_regress_multi_check = $(PERL) $(citus_abs_srcdir)/pg_regress_multi.pl --pgxsdir="$(pgxsdir)" --bindir="$(bindir)" --libdir="$(libdir)" --majorversion="$(MAJORVERSION)" --postgres-builddir="$(postgres_abs_builddir)" --postgres-srcdir="$(postgres_abs_srcdir)" --citus_abs_srcdir="$(citus_abs_srcdir)" --citus-version=$(CITUSVERSION)
+pg_regress_multi_check = $(PERL) $(citus_abs_srcdir)/pg_regress_multi.pl --pgxsdir="$(pgxsdir)" --bindir="$(bindir)" --libdir="$(libdir)" --majorversion="$(MAJORVERSION)" --postgres-builddir="$(postgres_abs_builddir)" --postgres-srcdir="$(postgres_abs_srcdir)" --citus_abs_srcdir="$(citus_abs_srcdir)" --citus-version=$(CITUSVERSION) --citus-libdir="$(CITUSLIBDIR)"
 MULTI_REGRESS_OPTS = --inputdir=$(citus_abs_srcdir) $(pg_regress_locale_flags) --launcher="$(citus_abs_srcdir)/log_test_times"
 
 pg_upgrade_check = $(citus_abs_srcdir)/citus_tests/upgrade/pg_upgrade_test.py
diff --git a/src/test/regress/pg_regress_multi.pl b/src/test/regress/pg_regress_multi.pl
index 7d760a2eb..bd2917ba0 100755
--- a/src/test/regress/pg_regress_multi.pl
+++ b/src/test/regress/pg_regress_multi.pl
@@ -51,6 +51,7 @@ sub Usage()
     print "  --mitmproxy        	Start a mitmproxy for one of the workers\n";
     print "  --worker-count         Number of workers in Citus cluster (default: 2)\n";
     print "  --citus-version        Citus version being tested (used for during extension create)\n";
+    print "  --citus-libdir         Absolute path to alternative Citus library directory\n";
     exit 1;
 }
 
@@ -89,6 +90,7 @@ my $publicWorker1Host = "localhost";
 my $publicWorker2Host = "localhost";
 my $workerCount = 2;
 my $citusversion = "";
+my $citusLibdir = "";
 
 my $serversAreShutdown = "TRUE";
 my $usingWindows = 0;
@@ -124,6 +126,7 @@ GetOptions(
     'worker-2-public-hostname=s' => \$publicWorker2Host,
     'worker-count=i' => \$workerCount,
     'citus-version=s' => \$citusversion,
+    'citus-libdir=s' => \$citusLibdir,
     'help' => sub { Usage() });
 
 my $fixopen = "$bindir/postgres.fixopen";
@@ -310,6 +313,40 @@ sub generate_hba
     close $fh;
 }
 
+sub setup_symlink
+{
+    my ($originalfile, $targetfile) = @_;
+
+    # Only proceed if not on Windows and both files are defined and non-empty
+    return if $usingWindows;
+    return unless (defined $originalfile && $originalfile ne "");
+    return unless (defined $targetfile && $targetfile ne "");
+
+    -e $targetfile or die "Target is not found at $targetfile";
+
+    my $backup = $originalfile . ".bak";
+    rename($originalfile, $backup) or die "Failed to copy $originalfile to $backup: $!";
+
+    symlink($targetfile, $originalfile) or die "Failed to create symlink $originalfile -> $targetfile: $!";
+}
+
+sub restore_original
+{
+    my ($originalfile) = @_;
+
+    # Only proceed if not on Windows and file is defined and non-empty
+    return if $usingWindows;
+    return unless (defined $originalfile && $originalfile ne "");
+
+    my $backup = $originalfile . ".bak";
+
+    # Return silently if backup doesn't exist
+    return unless -e $backup;
+
+    unlink($originalfile) or die "Failed to remove symlink at $originalfile: $!";
+    rename($backup, $originalfile) or die "Failed to restore original file from $backup to $originalfile: $!";
+}
+
 # always want to call initdb under normal postgres, so revert from a
 # partial run, even if we're now not using valgrind.
 revert_replace_postgres();
@@ -593,7 +630,7 @@ if($isolationtester)
    push(@pgOptions, "citus.background_task_queue_interval=-1");
 }
 
-if($citusversion)
+if($citusversion || $citusLibdir)
 {
     push(@pgOptions, "citus.enable_version_checks=off");
     push(@pgOptions, "columnar.enable_version_checks=off");
@@ -783,6 +820,9 @@ if (!$conninfo)
 # Routine to shutdown servers at failure/exit
 sub ShutdownServers()
 {
+    restore_original(catfile($libdir, "citus.so"));
+    restore_original(catfile($libdir, "citus_columnar.so"));
+
     if (!$conninfo && $serversAreShutdown eq "FALSE")
     {
         system(catfile("$bindir", "pg_ctl"),
@@ -908,6 +948,13 @@ if ($followercluster && $backupnodetest == 0)
     $synchronousReplication = "-c synchronous_standby_names='FIRST 1 (*)' -c synchronous_commit=remote_apply";
 }
 
+# Ensure citus.so points to alternative library if provided
+if ($citusLibdir)
+{
+    setup_symlink(catfile($libdir, "citus.so"), catfile($citusLibdir, "citus.so"));
+    setup_symlink(catfile($libdir, "citus_columnar.so"), catfile($citusLibdir, "citus_columnar.so"));
+}
+
 # Start servers
 if (!$conninfo)
 {
@@ -979,6 +1026,13 @@ if ($followercluster)
     }
 }
 
+# Restore original citus.so if we modified it
+if ($citusLibdir)
+{
+    restore_original(catfile($libdir, "citus.so"));
+    restore_original(catfile($libdir, "citus_columnar.so"));
+}
+
 ###
 # Create database, extensions, types, functions and fdws on the workers,
 # pg_regress won't know to create them for us.