From 287abea661b98ba1ebb72b88843ff5e4b41c0718 Mon Sep 17 00:00:00 2001
From: Naisila Puka <37271756+naisila@users.noreply.github.com>
Date: Thu, 9 Oct 2025 17:35:03 +0300
Subject: [PATCH] PG18 compatibility - varreturningtype additions (#8231)

This PR solves the following diffs, originating from the addition of
`varreturningtype` field to the `Var` struct in PG18:
https://github.com/postgres/postgres/commit/80feb727c

Previously we didn't account for this new field (as it's new), so this
wouldn't allow the parser to correctly reconstruct the `Var` node
structure, but rather it would error out with `did not find '}' at end
of input node`:

```diff
 SELECT column_to_column_name(logicalrelid, partkey)
 FROM pg_dist_partition WHERE partkey IS NOT NULL ORDER BY 1 LIMIT 1;
- column_to_column_name
----------------------------------------------------------------------
- a
-(1 row)
-
+ERROR:  did not find '}' at end of input node
```

Solution follows precedent https://github.com/citusdata/citus/pull/7107,
when varnullingrels field was added to the `Var` struct in PG16.

The solution includes:
- Taking care of the `partkey` in `pg_dist_partition` table because it's
coming from the `Var` struct. This mainly includes fixing the upgrade
script to PG18, by saving all the `partkey` infos before upgrading to
PG18 (in `citus_prepare_pg_upgrade`), and then re-generating `partkey`
columns in `pg_dist_partition` (using `UPDATE`) after upgrading to PG18
(in `citus_finish_pg_upgrade`).
- Adding a normalize rule to fix output differences among PG versions.
Note that we need two normalize lines: one for PG15 since it doesn't
have `varnullingrels`, and one for PG16/PG17.
- Small trick on `metadata_sync_helpers` to use different text when
generating the `partkey`, based on the PG version.

Fixes #8189
---
 .../distributed/sql/citus--13.2-1--14.0-1.sql |   3 +
 .../sql/downgrades/citus--14.0-1--13.2-1.sql  |   3 +
 .../udfs/citus_finish_pg_upgrade/14.0-1.sql   | 268 ++++++++++++++++++
 .../udfs/citus_finish_pg_upgrade/latest.sql   |   8 +
 .../udfs/citus_prepare_pg_upgrade/14.0-1.sql  | 111 ++++++++
 .../udfs/citus_prepare_pg_upgrade/latest.sql  |  11 +
 src/test/regress/bin/normalize.sed            |   8 +-
 .../regress/expected/citus_local_tables.out   |   2 +-
 .../expected/citus_local_tables_mx.out        |   4 +-
 .../expected/metadata_sync_helpers.out        |  10 +-
 src/test/regress/expected/multi_extension.out |   2 +-
 .../regress/expected/multi_metadata_sync.out  |   4 +-
 .../expected/multi_metadata_sync_0.out        |   4 +-
 .../expected/multi_prune_shard_list.out       |   2 +-
 .../expected/single_shard_table_udfs.out      |   2 +-
 .../expected/start_stop_metadata_sync.out     |  12 +-
 .../regress/expected/upgrade_basic_after.out  |   1 +
 .../regress/expected/upgrade_basic_before.out |   1 +
 .../regress/sql/metadata_sync_helpers.sql     |  10 +-
 src/test/regress/sql/multi_extension.sql      |   2 +-
 src/test/regress/sql/upgrade_basic_after.sql  |   1 +
 src/test/regress/sql/upgrade_basic_before.sql |   1 +
 22 files changed, 448 insertions(+), 22 deletions(-)
 create mode 100644 src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/14.0-1.sql
 create mode 100644 src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/14.0-1.sql

diff --git a/src/backend/distributed/sql/citus--13.2-1--14.0-1.sql b/src/backend/distributed/sql/citus--13.2-1--14.0-1.sql
index ebf34b7fd..815d4e794 100644
--- a/src/backend/distributed/sql/citus--13.2-1--14.0-1.sql
+++ b/src/backend/distributed/sql/citus--13.2-1--14.0-1.sql
@@ -1,2 +1,5 @@
 -- citus--13.2-1--14.0-1
 -- bump version to 14.0-1
+
+#include "udfs/citus_prepare_pg_upgrade/14.0-1.sql"
+#include "udfs/citus_finish_pg_upgrade/14.0-1.sql"
diff --git a/src/backend/distributed/sql/downgrades/citus--14.0-1--13.2-1.sql b/src/backend/distributed/sql/downgrades/citus--14.0-1--13.2-1.sql
index 3ef630c4b..199030339 100644
--- a/src/backend/distributed/sql/downgrades/citus--14.0-1--13.2-1.sql
+++ b/src/backend/distributed/sql/downgrades/citus--14.0-1--13.2-1.sql
@@ -1,2 +1,5 @@
 -- citus--14.0-1--13.2-1
 -- downgrade version to 13.2-1
+
+#include "../udfs/citus_prepare_pg_upgrade/13.0-1.sql"
+#include "../udfs/citus_finish_pg_upgrade/13.2-1.sql"
diff --git a/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/14.0-1.sql b/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/14.0-1.sql
new file mode 100644
index 000000000..4bad681eb
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/14.0-1.sql
@@ -0,0 +1,268 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_finish_pg_upgrade()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    AS $cppu$
+DECLARE
+    table_name regclass;
+    command text;
+    trigger_name text;
+BEGIN
+
+
+    IF substring(current_Setting('server_version'), '\d+')::int >= 14 THEN
+    EXECUTE $cmd$
+        -- disable propagation to prevent EnsureCoordinator errors
+        -- the aggregate created here does not depend on Citus extension (yet)
+        -- since we add the dependency with the next command
+        SET citus.enable_ddl_propagation TO OFF;
+        CREATE AGGREGATE array_cat_agg(anycompatiblearray) (SFUNC = array_cat, STYPE = anycompatiblearray);
+        COMMENT ON AGGREGATE array_cat_agg(anycompatiblearray)
+        IS 'concatenate input arrays into a single array';
+        RESET citus.enable_ddl_propagation;
+    $cmd$;
+    ELSE
+    EXECUTE $cmd$
+        SET citus.enable_ddl_propagation TO OFF;
+        CREATE AGGREGATE array_cat_agg(anyarray) (SFUNC = array_cat, STYPE = anyarray);
+        COMMENT ON AGGREGATE array_cat_agg(anyarray)
+        IS 'concatenate input arrays into a single array';
+        RESET citus.enable_ddl_propagation;
+    $cmd$;
+    END IF;
+
+    --
+    -- Citus creates the array_cat_agg but because of a compatibility
+    -- issue between pg13-pg14, we drop and create it during upgrade.
+    -- And as Citus creates it, there needs to be a dependency to the
+    -- Citus extension, so we create that dependency here.
+    -- We are not using:
+    --  ALTER EXENSION citus DROP/CREATE AGGREGATE array_cat_agg
+    -- because we don't have an easy way to check if the aggregate
+    -- exists with anyarray type or anycompatiblearray type.
+
+    INSERT INTO pg_depend
+    SELECT
+        'pg_proc'::regclass::oid as classid,
+        (SELECT oid FROM pg_proc WHERE proname = 'array_cat_agg') as objid,
+        0 as objsubid,
+        'pg_extension'::regclass::oid as refclassid,
+        (select oid from pg_extension where extname = 'citus') as refobjid,
+        0 as refobjsubid ,
+        'e' as deptype;
+
+    -- PG16 has its own any_value, so only create it pre PG16.
+    -- We can remove this part when we drop support for PG16
+    IF substring(current_Setting('server_version'), '\d+')::int < 16 THEN
+    EXECUTE $cmd$
+        -- disable propagation to prevent EnsureCoordinator errors
+        -- the aggregate created here does not depend on Citus extension (yet)
+        -- since we add the dependency with the next command
+        SET citus.enable_ddl_propagation TO OFF;
+        CREATE OR REPLACE FUNCTION pg_catalog.any_value_agg ( anyelement, anyelement )
+        RETURNS anyelement AS $$
+                SELECT CASE WHEN $1 IS NULL THEN $2 ELSE $1 END;
+        $$ LANGUAGE SQL STABLE;
+
+        CREATE AGGREGATE pg_catalog.any_value (
+                sfunc       = pg_catalog.any_value_agg,
+                combinefunc = pg_catalog.any_value_agg,
+                basetype    = anyelement,
+                stype       = anyelement
+        );
+        COMMENT ON AGGREGATE pg_catalog.any_value(anyelement) IS
+            'Returns the value of any row in the group. It is mostly useful when you know there will be only 1 element.';
+        RESET citus.enable_ddl_propagation;
+        --
+        -- Citus creates the any_value aggregate but because of a compatibility
+        -- issue between pg15-pg16 -- any_value is created in PG16, we drop
+        -- and create it during upgrade IF upgraded version is less than 16.
+        -- And as Citus creates it, there needs to be a dependency to the
+        -- Citus extension, so we create that dependency here.
+
+        INSERT INTO pg_depend
+        SELECT
+            'pg_proc'::regclass::oid as classid,
+            (SELECT oid FROM pg_proc WHERE proname = 'any_value_agg') as objid,
+            0 as objsubid,
+            'pg_extension'::regclass::oid as refclassid,
+            (select oid from pg_extension where extname = 'citus') as refobjid,
+            0 as refobjsubid ,
+            'e' as deptype;
+
+        INSERT INTO pg_depend
+        SELECT
+            'pg_proc'::regclass::oid as classid,
+            (SELECT oid FROM pg_proc WHERE proname = 'any_value') as objid,
+            0 as objsubid,
+            'pg_extension'::regclass::oid as refclassid,
+            (select oid from pg_extension where extname = 'citus') as refobjid,
+            0 as refobjsubid ,
+            'e' as deptype;
+    $cmd$;
+    END IF;
+
+    --
+    -- restore citus catalog tables
+    --
+    INSERT INTO pg_catalog.pg_dist_partition SELECT * FROM public.pg_dist_partition;
+
+    -- if we are upgrading from PG14/PG15 to PG16+,
+    -- we need to regenerate the partkeys because they will include varnullingrels as well.
+    UPDATE pg_catalog.pg_dist_partition
+    SET partkey = column_name_to_column(pg_dist_partkeys_pre_16_upgrade.logicalrelid, col_name)
+    FROM public.pg_dist_partkeys_pre_16_upgrade
+    WHERE pg_dist_partkeys_pre_16_upgrade.logicalrelid = pg_dist_partition.logicalrelid;
+    DROP TABLE public.pg_dist_partkeys_pre_16_upgrade;
+
+    -- if we are upgrading to PG18+,
+    -- we need to regenerate the partkeys because they will include varreturningtype as well.
+    UPDATE pg_catalog.pg_dist_partition
+    SET partkey = column_name_to_column(pg_dist_partkeys_pre_18_upgrade.logicalrelid, col_name)
+    FROM public.pg_dist_partkeys_pre_18_upgrade
+    WHERE pg_dist_partkeys_pre_18_upgrade.logicalrelid = pg_dist_partition.logicalrelid;
+    DROP TABLE public.pg_dist_partkeys_pre_18_upgrade;
+
+    INSERT INTO pg_catalog.pg_dist_shard SELECT * FROM public.pg_dist_shard;
+    INSERT INTO pg_catalog.pg_dist_placement SELECT * FROM public.pg_dist_placement;
+    INSERT INTO pg_catalog.pg_dist_node_metadata SELECT * FROM public.pg_dist_node_metadata;
+    INSERT INTO pg_catalog.pg_dist_node SELECT * FROM public.pg_dist_node;
+    INSERT INTO pg_catalog.pg_dist_local_group SELECT * FROM public.pg_dist_local_group;
+    INSERT INTO pg_catalog.pg_dist_transaction SELECT * FROM public.pg_dist_transaction;
+    INSERT INTO pg_catalog.pg_dist_colocation SELECT * FROM public.pg_dist_colocation;
+    INSERT INTO pg_catalog.pg_dist_cleanup SELECT * FROM public.pg_dist_cleanup;
+    INSERT INTO pg_catalog.pg_dist_schema SELECT schemaname::regnamespace, colocationid FROM public.pg_dist_schema;
+    -- enterprise catalog tables
+    INSERT INTO pg_catalog.pg_dist_authinfo SELECT * FROM public.pg_dist_authinfo;
+    INSERT INTO pg_catalog.pg_dist_poolinfo SELECT * FROM public.pg_dist_poolinfo;
+
+    -- Temporarily disable trigger to check for validity of functions while
+    -- inserting. The current contents of the table might be invalid if one of
+    -- the functions was removed by the user without also removing the
+    -- rebalance strategy. Obviously that's not great, but it should be no
+    -- reason to fail the upgrade.
+    ALTER TABLE pg_catalog.pg_dist_rebalance_strategy DISABLE TRIGGER pg_dist_rebalance_strategy_validation_trigger;
+    INSERT INTO pg_catalog.pg_dist_rebalance_strategy SELECT
+        name,
+        default_strategy,
+        shard_cost_function::regprocedure::regproc,
+        node_capacity_function::regprocedure::regproc,
+        shard_allowed_on_node_function::regprocedure::regproc,
+        default_threshold,
+        minimum_threshold,
+        improvement_threshold
+    FROM public.pg_dist_rebalance_strategy;
+    ALTER TABLE pg_catalog.pg_dist_rebalance_strategy ENABLE TRIGGER pg_dist_rebalance_strategy_validation_trigger;
+
+    --
+    -- drop backup tables
+    --
+    DROP TABLE public.pg_dist_authinfo;
+    DROP TABLE public.pg_dist_colocation;
+    DROP TABLE public.pg_dist_local_group;
+    DROP TABLE public.pg_dist_node;
+    DROP TABLE public.pg_dist_node_metadata;
+    DROP TABLE public.pg_dist_partition;
+    DROP TABLE public.pg_dist_placement;
+    DROP TABLE public.pg_dist_poolinfo;
+    DROP TABLE public.pg_dist_shard;
+    DROP TABLE public.pg_dist_transaction;
+    DROP TABLE public.pg_dist_rebalance_strategy;
+    DROP TABLE public.pg_dist_cleanup;
+    DROP TABLE public.pg_dist_schema;
+    --
+    -- reset sequences
+    --
+    PERFORM setval('pg_catalog.pg_dist_shardid_seq', (SELECT MAX(shardid)+1 AS max_shard_id FROM pg_dist_shard), false);
+    PERFORM setval('pg_catalog.pg_dist_placement_placementid_seq', (SELECT MAX(placementid)+1 AS max_placement_id FROM pg_dist_placement), false);
+    PERFORM setval('pg_catalog.pg_dist_groupid_seq', (SELECT MAX(groupid)+1 AS max_group_id FROM pg_dist_node), false);
+    PERFORM setval('pg_catalog.pg_dist_node_nodeid_seq', (SELECT MAX(nodeid)+1 AS max_node_id FROM pg_dist_node), false);
+    PERFORM setval('pg_catalog.pg_dist_colocationid_seq', (SELECT MAX(colocationid)+1 AS max_colocation_id FROM pg_dist_colocation), false);
+    PERFORM setval('pg_catalog.pg_dist_operationid_seq', (SELECT MAX(operation_id)+1 AS max_operation_id FROM pg_dist_cleanup), false);
+    PERFORM setval('pg_catalog.pg_dist_cleanup_recordid_seq', (SELECT MAX(record_id)+1 AS max_record_id FROM pg_dist_cleanup), false);
+    PERFORM setval('pg_catalog.pg_dist_clock_logical_seq', (SELECT last_value FROM public.pg_dist_clock_logical_seq), false);
+    DROP TABLE public.pg_dist_clock_logical_seq;
+
+
+
+    --
+    -- register triggers
+    --
+    FOR table_name IN SELECT logicalrelid FROM pg_catalog.pg_dist_partition JOIN pg_class ON (logicalrelid = oid) WHERE relkind <> 'f'
+    LOOP
+        trigger_name := 'truncate_trigger_' || table_name::oid;
+        command := 'create trigger ' || trigger_name || ' after truncate on ' || table_name || ' execute procedure pg_catalog.citus_truncate_trigger()';
+        EXECUTE command;
+        command := 'update pg_trigger set tgisinternal = true where tgname = ' || quote_literal(trigger_name);
+        EXECUTE command;
+    END LOOP;
+
+    --
+    -- set dependencies
+    --
+    INSERT INTO pg_depend
+    SELECT
+        'pg_class'::regclass::oid as classid,
+        p.logicalrelid::regclass::oid as objid,
+        0 as objsubid,
+        'pg_extension'::regclass::oid as refclassid,
+        (select oid from pg_extension where extname = 'citus') as refobjid,
+        0 as refobjsubid ,
+        'n' as deptype
+    FROM pg_catalog.pg_dist_partition p;
+
+    -- If citus_columnar extension exists, then perform the post PG-upgrade work for columnar as well.
+    --
+    -- First look if pg_catalog.columnar_finish_pg_upgrade function exists as part of the citus_columnar
+    -- extension. (We check whether it's part of the extension just for security reasons). If it does, then
+    -- call it. If not, then look for columnar_internal.columnar_ensure_am_depends_catalog function and as
+    -- part of the citus_columnar extension. If so, then call it. We alternatively check for the latter UDF
+    -- just because pg_catalog.columnar_finish_pg_upgrade function is introduced in citus_columnar 13.2-1
+    -- and as of today all it does is to call columnar_internal.columnar_ensure_am_depends_catalog function.
+    IF EXISTS (
+        SELECT 1 FROM pg_depend
+        JOIN pg_proc ON (pg_depend.objid = pg_proc.oid)
+        JOIN pg_namespace ON (pg_proc.pronamespace = pg_namespace.oid)
+        JOIN pg_extension ON (pg_depend.refobjid = pg_extension.oid)
+        WHERE
+            -- Looking if pg_catalog.columnar_finish_pg_upgrade function exists and
+            -- if there is a dependency record from it (proc class = 1255) ..
+            pg_depend.classid = 1255 AND pg_namespace.nspname = 'pg_catalog' AND pg_proc.proname = 'columnar_finish_pg_upgrade' AND
+            -- .. to citus_columnar extension (3079 = extension class), if it exists.
+            pg_depend.refclassid = 3079 AND pg_extension.extname = 'citus_columnar'
+    )
+    THEN PERFORM pg_catalog.columnar_finish_pg_upgrade();
+    ELSIF EXISTS (
+        SELECT 1 FROM pg_depend
+        JOIN pg_proc ON (pg_depend.objid = pg_proc.oid)
+        JOIN pg_namespace ON (pg_proc.pronamespace = pg_namespace.oid)
+        JOIN pg_extension ON (pg_depend.refobjid = pg_extension.oid)
+        WHERE
+            -- Looking if columnar_internal.columnar_ensure_am_depends_catalog function exists and
+            -- if there is a dependency record from it (proc class = 1255) ..
+            pg_depend.classid = 1255 AND pg_namespace.nspname = 'columnar_internal' AND pg_proc.proname = 'columnar_ensure_am_depends_catalog' AND
+            -- .. to citus_columnar extension (3079 = extension class), if it exists.
+            pg_depend.refclassid = 3079 AND pg_extension.extname = 'citus_columnar'
+    )
+    THEN PERFORM columnar_internal.columnar_ensure_am_depends_catalog();
+    END IF;
+
+    -- restore pg_dist_object from the stable identifiers
+    TRUNCATE pg_catalog.pg_dist_object;
+    INSERT INTO pg_catalog.pg_dist_object (classid, objid, objsubid, distribution_argument_index, colocationid)
+    SELECT
+        address.classid,
+        address.objid,
+        address.objsubid,
+        naming.distribution_argument_index,
+        naming.colocationid
+    FROM
+        public.pg_dist_object naming,
+        pg_catalog.pg_get_object_address(naming.type, naming.object_names, naming.object_args) address;
+
+    DROP TABLE public.pg_dist_object;
+END;
+$cppu$;
+
+COMMENT ON FUNCTION pg_catalog.citus_finish_pg_upgrade()
+    IS 'perform tasks to restore citus settings from a location that has been prepared before pg_upgrade';
diff --git a/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql b/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql
index 38daeb86c..4bad681eb 100644
--- a/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql
+++ b/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql
@@ -115,6 +115,14 @@ BEGIN
     WHERE pg_dist_partkeys_pre_16_upgrade.logicalrelid = pg_dist_partition.logicalrelid;
     DROP TABLE public.pg_dist_partkeys_pre_16_upgrade;
 
+    -- if we are upgrading to PG18+,
+    -- we need to regenerate the partkeys because they will include varreturningtype as well.
+    UPDATE pg_catalog.pg_dist_partition
+    SET partkey = column_name_to_column(pg_dist_partkeys_pre_18_upgrade.logicalrelid, col_name)
+    FROM public.pg_dist_partkeys_pre_18_upgrade
+    WHERE pg_dist_partkeys_pre_18_upgrade.logicalrelid = pg_dist_partition.logicalrelid;
+    DROP TABLE public.pg_dist_partkeys_pre_18_upgrade;
+
     INSERT INTO pg_catalog.pg_dist_shard SELECT * FROM public.pg_dist_shard;
     INSERT INTO pg_catalog.pg_dist_placement SELECT * FROM public.pg_dist_placement;
     INSERT INTO pg_catalog.pg_dist_node_metadata SELECT * FROM public.pg_dist_node_metadata;
diff --git a/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/14.0-1.sql b/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/14.0-1.sql
new file mode 100644
index 000000000..ee5e98e1d
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/14.0-1.sql
@@ -0,0 +1,111 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_prepare_pg_upgrade()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    AS $cppu$
+BEGIN
+
+    DELETE FROM pg_depend WHERE
+        objid IN (SELECT oid FROM pg_proc WHERE proname = 'array_cat_agg') AND
+        refobjid IN (select oid from pg_extension where extname = 'citus');
+    --
+    -- We are dropping the aggregates because postgres 14 changed
+    -- array_cat type from anyarray to anycompatiblearray. When
+    -- upgrading to pg14, specifically when running pg_restore on
+    -- array_cat_agg we would get an error. So we drop the aggregate
+    -- and create the right one on citus_finish_pg_upgrade.
+
+    DROP AGGREGATE IF EXISTS array_cat_agg(anyarray);
+    DROP AGGREGATE IF EXISTS array_cat_agg(anycompatiblearray);
+
+    -- We should drop any_value because PG16+ has its own any_value function
+    -- We can remove this part when we drop support for PG16
+    IF substring(current_Setting('server_version'), '\d+')::int < 16 THEN
+        DELETE FROM pg_depend WHERE
+            objid IN (SELECT oid FROM pg_proc WHERE proname = 'any_value' OR proname = 'any_value_agg') AND
+            refobjid IN (select oid from pg_extension where extname = 'citus');
+        DROP AGGREGATE IF EXISTS pg_catalog.any_value(anyelement);
+        DROP FUNCTION IF EXISTS pg_catalog.any_value_agg(anyelement, anyelement);
+    END IF;
+
+    --
+    -- Drop existing backup tables
+    --
+    DROP TABLE IF EXISTS public.pg_dist_partition;
+    DROP TABLE IF EXISTS public.pg_dist_shard;
+    DROP TABLE IF EXISTS public.pg_dist_placement;
+    DROP TABLE IF EXISTS public.pg_dist_node_metadata;
+    DROP TABLE IF EXISTS public.pg_dist_node;
+    DROP TABLE IF EXISTS public.pg_dist_local_group;
+    DROP TABLE IF EXISTS public.pg_dist_transaction;
+    DROP TABLE IF EXISTS public.pg_dist_colocation;
+    DROP TABLE IF EXISTS public.pg_dist_authinfo;
+    DROP TABLE IF EXISTS public.pg_dist_poolinfo;
+    DROP TABLE IF EXISTS public.pg_dist_rebalance_strategy;
+    DROP TABLE IF EXISTS public.pg_dist_object;
+    DROP TABLE IF EXISTS public.pg_dist_cleanup;
+    DROP TABLE IF EXISTS public.pg_dist_schema;
+    DROP TABLE IF EXISTS public.pg_dist_clock_logical_seq;
+
+    --
+    -- backup citus catalog tables
+    --
+    CREATE TABLE public.pg_dist_partition AS SELECT * FROM pg_catalog.pg_dist_partition;
+    CREATE TABLE public.pg_dist_shard AS SELECT * FROM pg_catalog.pg_dist_shard;
+    CREATE TABLE public.pg_dist_placement AS SELECT * FROM pg_catalog.pg_dist_placement;
+    CREATE TABLE public.pg_dist_node_metadata AS SELECT * FROM pg_catalog.pg_dist_node_metadata;
+    CREATE TABLE public.pg_dist_node AS SELECT * FROM pg_catalog.pg_dist_node;
+    CREATE TABLE public.pg_dist_local_group AS SELECT * FROM pg_catalog.pg_dist_local_group;
+    CREATE TABLE public.pg_dist_transaction AS SELECT * FROM pg_catalog.pg_dist_transaction;
+    CREATE TABLE public.pg_dist_colocation AS SELECT * FROM pg_catalog.pg_dist_colocation;
+    CREATE TABLE public.pg_dist_cleanup AS SELECT * FROM pg_catalog.pg_dist_cleanup;
+    -- save names of the tenant schemas instead of their oids because the oids might change after pg upgrade
+    CREATE TABLE public.pg_dist_schema AS SELECT schemaid::regnamespace::text AS schemaname, colocationid FROM pg_catalog.pg_dist_schema;
+    -- enterprise catalog tables
+    CREATE TABLE public.pg_dist_authinfo AS SELECT * FROM pg_catalog.pg_dist_authinfo;
+    CREATE TABLE public.pg_dist_poolinfo AS SELECT * FROM pg_catalog.pg_dist_poolinfo;
+    -- sequences
+    CREATE TABLE public.pg_dist_clock_logical_seq AS SELECT last_value FROM pg_catalog.pg_dist_clock_logical_seq;
+    CREATE TABLE public.pg_dist_rebalance_strategy AS SELECT
+        name,
+        default_strategy,
+        shard_cost_function::regprocedure::text,
+        node_capacity_function::regprocedure::text,
+        shard_allowed_on_node_function::regprocedure::text,
+        default_threshold,
+        minimum_threshold,
+        improvement_threshold
+    FROM pg_catalog.pg_dist_rebalance_strategy;
+
+    -- store upgrade stable identifiers on pg_dist_object catalog
+    CREATE TABLE public.pg_dist_object AS SELECT
+       address.type,
+       address.object_names,
+       address.object_args,
+       objects.distribution_argument_index,
+       objects.colocationid
+    FROM pg_catalog.pg_dist_object objects,
+         pg_catalog.pg_identify_object_as_address(objects.classid, objects.objid, objects.objsubid) address;
+
+    -- if we are upgrading from PG14/PG15 to PG16+,
+    -- we will need to regenerate the partkeys because they will include varnullingrels as well.
+    -- so we save the partkeys as column names here
+    CREATE TABLE IF NOT EXISTS public.pg_dist_partkeys_pre_16_upgrade AS
+    SELECT logicalrelid, column_to_column_name(logicalrelid, partkey) as col_name
+    FROM pg_catalog.pg_dist_partition WHERE partkey IS NOT NULL AND partkey NOT ILIKE '%varnullingrels%';
+
+    -- similarly, if we are upgrading to PG18+,
+    -- we will need to regenerate the partkeys because they will include varreturningtype as well.
+    -- so we save the partkeys as column names here
+    CREATE TABLE IF NOT EXISTS public.pg_dist_partkeys_pre_18_upgrade AS
+    SELECT logicalrelid, column_to_column_name(logicalrelid, partkey) as col_name
+    FROM pg_catalog.pg_dist_partition WHERE partkey IS NOT NULL AND partkey NOT ILIKE '%varreturningtype%';
+    -- remove duplicates (we would only have duplicates if we are upgrading from pre-16 to PG18+)
+    DELETE FROM public.pg_dist_partkeys_pre_18_upgrade USING public.pg_dist_partkeys_pre_16_upgrade p16
+    WHERE public.pg_dist_partkeys_pre_18_upgrade.logicalrelid = p16.logicalrelid
+    AND public.pg_dist_partkeys_pre_18_upgrade.col_name = p16.col_name;
+END;
+$cppu$;
+
+COMMENT ON FUNCTION pg_catalog.citus_prepare_pg_upgrade()
+    IS 'perform tasks to copy citus settings to a location that could later be restored after pg_upgrade is done';
diff --git a/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql b/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql
index 4f07ce5c4..ee5e98e1d 100644
--- a/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql
+++ b/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql
@@ -93,6 +93,17 @@ BEGIN
     CREATE TABLE IF NOT EXISTS public.pg_dist_partkeys_pre_16_upgrade AS
     SELECT logicalrelid, column_to_column_name(logicalrelid, partkey) as col_name
     FROM pg_catalog.pg_dist_partition WHERE partkey IS NOT NULL AND partkey NOT ILIKE '%varnullingrels%';
+
+    -- similarly, if we are upgrading to PG18+,
+    -- we will need to regenerate the partkeys because they will include varreturningtype as well.
+    -- so we save the partkeys as column names here
+    CREATE TABLE IF NOT EXISTS public.pg_dist_partkeys_pre_18_upgrade AS
+    SELECT logicalrelid, column_to_column_name(logicalrelid, partkey) as col_name
+    FROM pg_catalog.pg_dist_partition WHERE partkey IS NOT NULL AND partkey NOT ILIKE '%varreturningtype%';
+    -- remove duplicates (we would only have duplicates if we are upgrading from pre-16 to PG18+)
+    DELETE FROM public.pg_dist_partkeys_pre_18_upgrade USING public.pg_dist_partkeys_pre_16_upgrade p16
+    WHERE public.pg_dist_partkeys_pre_18_upgrade.logicalrelid = p16.logicalrelid
+    AND public.pg_dist_partkeys_pre_18_upgrade.col_name = p16.col_name;
 END;
 $cppu$;
 
diff --git a/src/test/regress/bin/normalize.sed b/src/test/regress/bin/normalize.sed
index 1ebf4392e..efdff159d 100644
--- a/src/test/regress/bin/normalize.sed
+++ b/src/test/regress/bin/normalize.sed
@@ -297,7 +297,13 @@ s/(NOTICE:  issuing CREATE EXTENSION IF NOT EXISTS citus_columnar WITH SCHEMA  p
 
 s/, password_required=false//g
 s/provide the file or change sslmode/provide the file, use the system's trusted roots with sslrootcert=system, or change sslmode/g
-s/(:varcollid [0-9]+) :varlevelsup 0/\1 :varnullingrels (b) :varlevelsup 0/g
+
+#pg18 varreturningtype - change needed for PG16, PG17 tests
+s/(:varnullingrels \(b\) :varlevelsup 0) (:varnosyn 1)/\1 :varreturningtype 0 \2/g
+
+#pg16 varnullingrels and pg18 varreturningtype - change needed for PG15 tests
+s/(:varcollid [0-9]+) :varlevelsup 0/\1 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0/g
+
 s/table_name_for_view\.([_a-z0-9]+)(,| |$)/\1\2/g
 s/permission denied to terminate process/must be a superuser to terminate superuser process/g
 s/permission denied to cancel query/must be a superuser to cancel superuser query/g
diff --git a/src/test/regress/expected/citus_local_tables.out b/src/test/regress/expected/citus_local_tables.out
index 4f3053094..cdec0aafa 100644
--- a/src/test/regress/expected/citus_local_tables.out
+++ b/src/test/regress/expected/citus_local_tables.out
@@ -635,7 +635,7 @@ FROM pg_dist_partition WHERE logicalrelid = 'citus_local_table_4'::regclass;
 SELECT column_name_to_column('citus_local_table_4', 'a');
                                                   column_name_to_column
 ---------------------------------------------------------------------
- {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1}
+ {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1}
 (1 row)
 
 SELECT master_update_shard_statistics(shardid)
diff --git a/src/test/regress/expected/citus_local_tables_mx.out b/src/test/regress/expected/citus_local_tables_mx.out
index 8b3cb953f..98c03fea3 100644
--- a/src/test/regress/expected/citus_local_tables_mx.out
+++ b/src/test/regress/expected/citus_local_tables_mx.out
@@ -769,8 +769,8 @@ SELECT logicalrelid, partmethod, partkey FROM pg_dist_partition
         ORDER BY logicalrelid;
      logicalrelid     | partmethod |                                                          partkey
 ---------------------------------------------------------------------
- parent_dropped_col   | h          | {VAR :varno 1 :varattno 1 :vartype 1082 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1}
- parent_dropped_col_2 | h          | {VAR :varno 1 :varattno 5 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 5 :location -1}
+ parent_dropped_col   | h          | {VAR :varno 1 :varattno 1 :vartype 1082 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1}
+ parent_dropped_col_2 | h          | {VAR :varno 1 :varattno 5 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 5 :location -1}
 (2 rows)
 
 -- some tests for view propagation on citus local tables
diff --git a/src/test/regress/expected/metadata_sync_helpers.out b/src/test/regress/expected/metadata_sync_helpers.out
index 9db68eaf5..6345e15ac 100644
--- a/src/test/regress/expected/metadata_sync_helpers.out
+++ b/src/test/regress/expected/metadata_sync_helpers.out
@@ -1287,11 +1287,17 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	SET application_name to 'citus_internal gpid=10000000001';
 	-- with an ugly trick, update the vartype of table from int to bigint
 	-- so that making two tables colocated fails
-	-- include varnullingrels for PG16
+	-- include varnullingrels for PG16+
     SHOW server_version \gset
     SELECT substring(:'server_version', '\d+')::int >= 16 AS server_version_ge_16
     \gset
-    \if :server_version_ge_16
+	-- include varreturningtype for PG18+
+	SELECT substring(:'server_version', '\d+')::int >= 18 AS server_version_ge_18
+    \gset
+	\if :server_version_ge_18
+	UPDATE pg_dist_partition SET partkey = '{VAR :varno 1 :varattno 1 :vartype 20 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 1 :varreturningtype 0 :varnoold 1 :varoattno 1 :location -1}'
+	WHERE logicalrelid = 'test_2'::regclass;
+    \elif :server_version_ge_16
     UPDATE pg_dist_partition SET partkey = '{VAR :varno 1 :varattno 1 :vartype 20 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 1 :varnoold 1 :varoattno 1 :location -1}'
 	WHERE logicalrelid = 'test_2'::regclass;
     \else
diff --git a/src/test/regress/expected/multi_extension.out b/src/test/regress/expected/multi_extension.out
index ce6a053ff..99774d9ef 100644
--- a/src/test/regress/expected/multi_extension.out
+++ b/src/test/regress/expected/multi_extension.out
@@ -1130,7 +1130,7 @@ DELETE FROM pg_dist_shard WHERE shardid = 1;
 CREATE TABLE e_transactions(order_id varchar(255) NULL, transaction_id int) PARTITION BY LIST(transaction_id);
 CREATE TABLE orders_2020_07_01
 PARTITION OF e_transactions FOR VALUES IN (1,2,3);
-INSERT INTO pg_dist_partition VALUES ('e_transactions'::regclass,'h', '{VAR :varno 1 :varattno 1 :vartype 1043 :vartypmod 259 :varcollid 100 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1}', 7, 's');
+INSERT INTO pg_dist_partition VALUES ('e_transactions'::regclass,'h', '{VAR :varno 1 :varattno 1 :vartype 1043 :vartypmod 259 :varcollid 100 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1}', 7, 's');
 SELECT
 	(metadata->>'partitioned_citus_table_exists_pre_11')::boolean as partitioned_citus_table_exists_pre_11,
 	(metadata->>'partitioned_citus_table_exists_pre_11') IS NULL as is_null
diff --git a/src/test/regress/expected/multi_metadata_sync.out b/src/test/regress/expected/multi_metadata_sync.out
index a68d8c03d..e9b5b587e 100644
--- a/src/test/regress/expected/multi_metadata_sync.out
+++ b/src/test/regress/expected/multi_metadata_sync.out
@@ -525,7 +525,7 @@ SELECT * FROM pg_dist_node ORDER BY nodeid;
 SELECT * FROM pg_dist_partition WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY logicalrelid::text;
           logicalrelid           | partmethod |                                                         partkey                                                          | colocationid | repmodel | autoconverted
 ---------------------------------------------------------------------
- mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
+ mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
 (1 row)
 
 SELECT * FROM pg_dist_shard  WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY shardid;
@@ -664,7 +664,7 @@ SELECT * FROM pg_dist_node ORDER BY nodeid;
 SELECT * FROM pg_dist_partition WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY logicalrelid::text;
           logicalrelid           | partmethod |                                                         partkey                                                          | colocationid | repmodel | autoconverted
 ---------------------------------------------------------------------
- mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
+ mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
 (1 row)
 
 SELECT * FROM pg_dist_shard WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY shardid;
diff --git a/src/test/regress/expected/multi_metadata_sync_0.out b/src/test/regress/expected/multi_metadata_sync_0.out
index f4a89ab60..5ff926ff6 100644
--- a/src/test/regress/expected/multi_metadata_sync_0.out
+++ b/src/test/regress/expected/multi_metadata_sync_0.out
@@ -535,7 +535,7 @@ SELECT * FROM pg_dist_node ORDER BY nodeid;
 SELECT * FROM pg_dist_partition WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY logicalrelid::text;
           logicalrelid           | partmethod |                                                         partkey                                                          | colocationid | repmodel | autoconverted
 ---------------------------------------------------------------------
- mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
+ mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
 (1 row)
 
 SELECT * FROM pg_dist_shard  WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY shardid;
@@ -674,7 +674,7 @@ SELECT * FROM pg_dist_node ORDER BY nodeid;
 SELECT * FROM pg_dist_partition WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY logicalrelid::text;
           logicalrelid           | partmethod |                                                         partkey                                                          | colocationid | repmodel | autoconverted
 ---------------------------------------------------------------------
- mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
+ mx_testing_schema.mx_test_table | h          | {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |            2 | s        | f
 (1 row)
 
 SELECT * FROM pg_dist_shard WHERE logicalrelid::text LIKE 'mx_testing_schema%' ORDER BY shardid;
diff --git a/src/test/regress/expected/multi_prune_shard_list.out b/src/test/regress/expected/multi_prune_shard_list.out
index 3762bd05d..88c385a46 100644
--- a/src/test/regress/expected/multi_prune_shard_list.out
+++ b/src/test/regress/expected/multi_prune_shard_list.out
@@ -86,7 +86,7 @@ SELECT prune_using_both_values('pruning', 'tomato', 'rose');
 SELECT debug_equality_expression('pruning');
                                                                                                                                                                             debug_equality_expression
 ---------------------------------------------------------------------
- {OPEXPR :opno 98 :opfuncid 67 :opresulttype 16 :opretset false :opcollid 0 :inputcollid 100 :args ({VAR :varno 1 :varattno 1 :vartype 25 :vartypmod -1 :varcollid 100 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} {CONST :consttype 25 :consttypmod -1 :constcollid 100 :constlen -1 :constbyval false :constisnull true :location -1 :constvalue <>}) :location -1}
+ {OPEXPR :opno 98 :opfuncid 67 :opresulttype 16 :opretset false :opcollid 0 :inputcollid 100 :args ({VAR :varno 1 :varattno 1 :vartype 25 :vartypmod -1 :varcollid 100 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} {CONST :consttype 25 :consttypmod -1 :constcollid 100 :constlen -1 :constbyval false :constisnull true :location -1 :constvalue <>}) :location -1}
 (1 row)
 
 -- print the initial ordering of shard intervals
diff --git a/src/test/regress/expected/single_shard_table_udfs.out b/src/test/regress/expected/single_shard_table_udfs.out
index 26b3f82cc..4bcd1185c 100644
--- a/src/test/regress/expected/single_shard_table_udfs.out
+++ b/src/test/regress/expected/single_shard_table_udfs.out
@@ -121,7 +121,7 @@ FROM pg_dist_partition WHERE logicalrelid = 'null_dist_key_table'::regclass;
 SELECT column_name_to_column('null_dist_key_table', 'a');
                                                   column_name_to_column
 ---------------------------------------------------------------------
- {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1}
+ {VAR :varno 1 :varattno 1 :vartype 23 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1}
 (1 row)
 
 SELECT master_update_shard_statistics(shardid)
diff --git a/src/test/regress/expected/start_stop_metadata_sync.out b/src/test/regress/expected/start_stop_metadata_sync.out
index ec9b0a034..c031700b3 100644
--- a/src/test/regress/expected/start_stop_metadata_sync.out
+++ b/src/test/regress/expected/start_stop_metadata_sync.out
@@ -158,12 +158,12 @@ SELECT * FROM test_matview;
 SELECT * FROM pg_dist_partition WHERE logicalrelid::text LIKE 'events%' ORDER BY logicalrelid::text;
         logicalrelid        | partmethod |                                                          partkey                                                           | colocationid | repmodel | autoconverted
 ---------------------------------------------------------------------
- events                     | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980000 | s        | f
- events_2021_feb            | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980000 | s        | f
- events_2021_jan            | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980000 | s        | f
- events_replicated          | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980001 | c        | f
- events_replicated_2021_feb | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980001 | c        | f
- events_replicated_2021_jan | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980001 | c        | f
+ events                     | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980000 | s        | f
+ events_2021_feb            | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980000 | s        | f
+ events_2021_jan            | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980000 | s        | f
+ events_replicated          | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980001 | c        | f
+ events_replicated_2021_feb | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980001 | c        | f
+ events_replicated_2021_jan | h          | {VAR :varno 1 :varattno 1 :vartype 1184 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1} |       980001 | c        | f
 (6 rows)
 
 SELECT count(*) > 0 FROM pg_dist_node;
diff --git a/src/test/regress/expected/upgrade_basic_after.out b/src/test/regress/expected/upgrade_basic_after.out
index 7c0ebfb29..2ce25bc54 100644
--- a/src/test/regress/expected/upgrade_basic_after.out
+++ b/src/test/regress/expected/upgrade_basic_after.out
@@ -307,6 +307,7 @@ SELECT * FROM t_range ORDER BY id;
 
 ROLLBACK;
 -- There is a difference in partkey Var representation between PG16 and older versions
+-- There is also a difference between PG18 and older versions
 -- Sanity check here that we can properly do column_to_column_name
 SELECT column_to_column_name(logicalrelid, partkey)
 FROM pg_dist_partition WHERE partkey IS NOT NULL ORDER BY 1 LIMIT 1;
diff --git a/src/test/regress/expected/upgrade_basic_before.out b/src/test/regress/expected/upgrade_basic_before.out
index 9abb6c806..99c523bb0 100644
--- a/src/test/regress/expected/upgrade_basic_before.out
+++ b/src/test/regress/expected/upgrade_basic_before.out
@@ -66,6 +66,7 @@ UPDATE pg_dist_shard SET shardminvalue = '5', shardmaxvalue = '7' WHERE shardid
 \copy t_range FROM STDIN with (DELIMITER ',')
 \copy t_range FROM STDIN with (DELIMITER ',')
 -- There is a difference in partkey Var representation between PG16 and older versions
+-- There is also a difference between PG18 and older versions
 -- Sanity check here that we can properly do column_to_column_name
 SELECT column_to_column_name(logicalrelid, partkey)
 FROM pg_dist_partition WHERE partkey IS NOT NULL ORDER BY 1 LIMIT 1;
diff --git a/src/test/regress/sql/metadata_sync_helpers.sql b/src/test/regress/sql/metadata_sync_helpers.sql
index dae331d25..9e67fa337 100644
--- a/src/test/regress/sql/metadata_sync_helpers.sql
+++ b/src/test/regress/sql/metadata_sync_helpers.sql
@@ -798,11 +798,17 @@ BEGIN TRANSACTION ISOLATION LEVEL READ COMMITTED;
 	-- with an ugly trick, update the vartype of table from int to bigint
 	-- so that making two tables colocated fails
 
-    -- include varnullingrels for PG16
+    -- include varnullingrels for PG16+
     SHOW server_version \gset
     SELECT substring(:'server_version', '\d+')::int >= 16 AS server_version_ge_16
     \gset
-    \if :server_version_ge_16
+	-- include varreturningtype for PG18+
+	SELECT substring(:'server_version', '\d+')::int >= 18 AS server_version_ge_18
+    \gset
+	\if :server_version_ge_18
+	UPDATE pg_dist_partition SET partkey = '{VAR :varno 1 :varattno 1 :vartype 20 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 1 :varreturningtype 0 :varnoold 1 :varoattno 1 :location -1}'
+	WHERE logicalrelid = 'test_2'::regclass;
+    \elif :server_version_ge_16
     UPDATE pg_dist_partition SET partkey = '{VAR :varno 1 :varattno 1 :vartype 20 :vartypmod -1 :varcollid 0 :varnullingrels (b) :varlevelsup 1 :varnoold 1 :varoattno 1 :location -1}'
 	WHERE logicalrelid = 'test_2'::regclass;
     \else
diff --git a/src/test/regress/sql/multi_extension.sql b/src/test/regress/sql/multi_extension.sql
index b7d69255b..90d565966 100644
--- a/src/test/regress/sql/multi_extension.sql
+++ b/src/test/regress/sql/multi_extension.sql
@@ -544,7 +544,7 @@ DELETE FROM pg_dist_shard WHERE shardid = 1;
 CREATE TABLE e_transactions(order_id varchar(255) NULL, transaction_id int) PARTITION BY LIST(transaction_id);
 CREATE TABLE orders_2020_07_01
 PARTITION OF e_transactions FOR VALUES IN (1,2,3);
-INSERT INTO pg_dist_partition VALUES ('e_transactions'::regclass,'h', '{VAR :varno 1 :varattno 1 :vartype 1043 :vartypmod 259 :varcollid 100 :varnullingrels (b) :varlevelsup 0 :varnosyn 1 :varattnosyn 1 :location -1}', 7, 's');
+INSERT INTO pg_dist_partition VALUES ('e_transactions'::regclass,'h', '{VAR :varno 1 :varattno 1 :vartype 1043 :vartypmod 259 :varcollid 100 :varnullingrels (b) :varlevelsup 0 :varreturningtype 0 :varnosyn 1 :varattnosyn 1 :location -1}', 7, 's');
 
 SELECT
 	(metadata->>'partitioned_citus_table_exists_pre_11')::boolean as partitioned_citus_table_exists_pre_11,
diff --git a/src/test/regress/sql/upgrade_basic_after.sql b/src/test/regress/sql/upgrade_basic_after.sql
index 855c06008..10604f6c3 100644
--- a/src/test/regress/sql/upgrade_basic_after.sql
+++ b/src/test/regress/sql/upgrade_basic_after.sql
@@ -103,6 +103,7 @@ SELECT * FROM t_range ORDER BY id;
 ROLLBACK;
 
 -- There is a difference in partkey Var representation between PG16 and older versions
+-- There is also a difference between PG18 and older versions
 -- Sanity check here that we can properly do column_to_column_name
 SELECT column_to_column_name(logicalrelid, partkey)
 FROM pg_dist_partition WHERE partkey IS NOT NULL ORDER BY 1 LIMIT 1;
diff --git a/src/test/regress/sql/upgrade_basic_before.sql b/src/test/regress/sql/upgrade_basic_before.sql
index 868483264..e81544d2a 100644
--- a/src/test/regress/sql/upgrade_basic_before.sql
+++ b/src/test/regress/sql/upgrade_basic_before.sql
@@ -53,6 +53,7 @@ UPDATE pg_dist_shard SET shardminvalue = '5', shardmaxvalue = '7' WHERE shardid
 \.
 
 -- There is a difference in partkey Var representation between PG16 and older versions
+-- There is also a difference between PG18 and older versions
 -- Sanity check here that we can properly do column_to_column_name
 SELECT column_to_column_name(logicalrelid, partkey)
 FROM pg_dist_partition WHERE partkey IS NOT NULL ORDER BY 1 LIMIT 1;