From 15a056aa2efc68cf632f75d462ff613348c6bf53 Mon Sep 17 00:00:00 2001
From: Staticle <55652600+staticle@users.noreply.github.com>
Date: Wed, 19 Feb 2025 12:53:42 +0000
Subject: [PATCH] Add Citus Support

---
 postgres-appliance/Dockerfile               |  7 +++-
 postgres-appliance/build_scripts/base.sh    | 36 +++++++++++++++++++++
 postgres-appliance/scripts/spilo_commons.py | 25 ++++++++++++--
 3 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/postgres-appliance/Dockerfile b/postgres-appliance/Dockerfile
index 9844872ea..1f376b6d7 100644
--- a/postgres-appliance/Dockerfile
+++ b/postgres-appliance/Dockerfile
@@ -60,7 +60,12 @@ ENV POSTGIS_VERSION=3.5 \
     PLPROFILER=REL4_2_5 \
     PG_PROFILE=4.7 \
     PAM_OAUTH2=v1.0.1 \
-    PG_PERMISSIONS_COMMIT=f4b7c18676fa64236a1c8e28d34a35764e4a70e2
+    PG_PERMISSIONS_COMMIT=f4b7c18676fa64236a1c8e28d34a35764e4a70e2 \
+    CITUS_PG_17_VERSION=13.0 \
+    CITUS_PG_16_VERSION=13.0 \
+    CITUS_PG_15_VERSION=13.0 \
+    CITUS_PG_14_VERSION=12.1 \
+    CITUS_PG_13_VERSION=11.3
 
 WORKDIR /builddeps
 RUN bash base.sh
diff --git a/postgres-appliance/build_scripts/base.sh b/postgres-appliance/build_scripts/base.sh
index f737f3652..f14258a92 100644
--- a/postgres-appliance/build_scripts/base.sh
+++ b/postgres-appliance/build_scripts/base.sh
@@ -176,6 +176,42 @@ if [ "$DEMO" != "true" ]; then
     done
 fi
 
+# install citus
+
+if [ "$(dpkg --print-architecture)" = "arm64" ]; then
+    echo "Citus was not installed as arm64 is not supported"
+else
+    curl https://install.citusdata.com/community/deb.sh | bash
+fi
+            
+for version in $DEB_PG_SUPPORTED_VERSIONS; do
+    CITUS_VERSION=""
+    case $version in
+        17)
+            CITUS_VERSION="$CITUS_PG_17_VERSION"
+            ;;
+        16)
+            CITUS_VERSION="$CITUS_PG_16_VERSION"
+            ;;
+        15)
+            CITUS_VERSION="$CITUS_PG_15_VERSION"
+            ;;
+        14)
+            CITUS_VERSION="$CITUS_PG_14_VERSION"
+            ;;
+        13)
+            CITUS_VERSION="$CITUS_PG_13_VERSION"
+            ;;
+        *)
+            ;;
+    esac
+    if [ "$CITUS_VERSION" != "" ] && [ "$(dpkg --print-architecture)" != "arm64" ]; then
+        apt-get -y install postgresql-"$version"-citus-"$CITUS_VERSION";
+    else
+        echo "Citus was not installed as appropriate version for Postgres $version amd64 was not provided";
+    fi
+done
+
 # make it possible for cron to work without root
 gcc -s -shared -fPIC -o /usr/local/lib/cron_unprivileged.so cron_unprivileged.c
 
diff --git a/postgres-appliance/scripts/spilo_commons.py b/postgres-appliance/scripts/spilo_commons.py
index 0543bf771..61a38e7de 100644
--- a/postgres-appliance/scripts/spilo_commons.py
+++ b/postgres-appliance/scripts/spilo_commons.py
@@ -15,19 +15,34 @@
     'timescaledb':    (9.6, 17, True,  True),
     'pg_cron':        (9.5, 17, True,  False),
     'pg_stat_kcache': (9.4, 17, True,  False),
-    'pg_partman':     (9.4, 17, False, True)
+    'pg_partman':     (9.4, 17, False, True),
+    'citus':          (13, 17, True, False)
 }
 if os.environ.get('ENABLE_PG_MON') == 'true':
     extensions['pg_mon'] = (11,  17, True,  False)
 
 
+def capture_shell_command_output(command):
+    result = subprocess.run(command, capture_output=True, text=True, shell=True)
+    output = result.stdout.strip()
+    return output
+
+
+def check_if_amd64_arch():
+    return capture_shell_command_output("dpkg --print-architecture") == "amd64"
+
+
 def adjust_extensions(old, version, extwlist=False):
     ret = []
     for name in old.split(','):
         name = name.strip()
         value = extensions.get(name)
         if name not in ret and value is None or value[0] <= version <= value[1] and (not extwlist or value[3]):
-            ret.append(name)
+            if name == "citus" and check_if_amd64_arch():
+                # Put Citus in first place in hared_preload_libraries
+                ret.insert(0, name)
+            else:
+                ret.append(name)
     return ','.join(ret)
 
 
@@ -38,7 +53,11 @@ def append_extensions(old, version, extwlist=False):
     def maybe_append(name):
         value = extensions.get(name)
         if name not in ret and (value is None or value[0] <= version <= value[1] and value[extwlist]):
-            ret.append(name)
+            if name == "citus" and check_if_amd64_arch():
+                # Put Citus in first place in hared_preload_libraries
+                ret.insert(0, name)
+            else:
+                ret.append(name)
 
     for name in old.split(','):
         maybe_append(name.strip())