-- citus--13.2-1--14.0-1
-- bump version to 14.0-1
CREATE OR REPLACE FUNCTION pg_catalog.citus_prepare_pg_upgrade()
    RETURNS void
    LANGUAGE plpgsql
    SET search_path = pg_catalog, pg_temp
    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';
CREATE OR REPLACE FUNCTION pg_catalog.citus_finish_pg_upgrade()
    RETURNS void
    LANGUAGE plpgsql
    SET search_path = pg_catalog, pg_temp
    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';


CREATE FUNCTION pg_catalog.worker_binary_partial_agg_ffunc(internal)
RETURNS bytea
AS 'MODULE_PATHNAME'
LANGUAGE C PARALLEL SAFE;
COMMENT ON FUNCTION pg_catalog.worker_binary_partial_agg_ffunc(internal)
    IS 'finalizer for worker_binary_partial_agg';
REVOKE ALL ON FUNCTION pg_catalog.worker_binary_partial_agg_ffunc FROM PUBLIC;
GRANT EXECUTE ON FUNCTION pg_catalog.worker_binary_partial_agg_ffunc TO PUBLIC;

CREATE FUNCTION pg_catalog.coord_binary_combine_agg_sfunc(internal, oid, bytea, anyelement)
RETURNS internal
AS 'MODULE_PATHNAME'
LANGUAGE C PARALLEL SAFE;
COMMENT ON FUNCTION pg_catalog.coord_binary_combine_agg_sfunc(internal, oid, bytea, anyelement)
    IS 'transition function for coord_binary_combine_agg';
REVOKE ALL ON FUNCTION pg_catalog.coord_binary_combine_agg_sfunc FROM PUBLIC;
GRANT EXECUTE ON FUNCTION pg_catalog.coord_binary_combine_agg_sfunc TO PUBLIC;

CREATE FUNCTION pg_catalog.coord_binary_combine_agg_ffunc(internal, oid, bytea, anyelement)
RETURNS anyelement
AS 'MODULE_PATHNAME'
LANGUAGE C PARALLEL SAFE;
COMMENT ON FUNCTION pg_catalog.coord_binary_combine_agg_ffunc(internal, oid, bytea, anyelement)
    IS 'finalizer for coord_binary_combine_agg';
REVOKE ALL ON FUNCTION pg_catalog.coord_binary_combine_agg_ffunc FROM PUBLIC;
GRANT EXECUTE ON FUNCTION pg_catalog.coord_binary_combine_agg_ffunc TO PUBLIC;

-- similar to worker_partial_agg but returns binary representation of the state
CREATE AGGREGATE pg_catalog.worker_binary_partial_agg(oid, anyelement) (
    STYPE = internal,
    SFUNC = pg_catalog.worker_partial_agg_sfunc,
    FINALFUNC = pg_catalog.worker_binary_partial_agg_ffunc
);
COMMENT ON AGGREGATE pg_catalog.worker_binary_partial_agg(oid, anyelement)
    IS 'support aggregate for implementing partial binary aggregation on workers';
REVOKE ALL ON FUNCTION pg_catalog.worker_binary_partial_agg FROM PUBLIC;
GRANT EXECUTE ON FUNCTION pg_catalog.worker_binary_partial_agg TO PUBLIC;

-- select coord_binary_combine_agg(agg, col) is similar to coord_combine_agg but
-- takes binary representation of the state as input
CREATE AGGREGATE pg_catalog.coord_binary_combine_agg(oid, bytea, anyelement) (
    STYPE = internal,
    SFUNC = pg_catalog.coord_binary_combine_agg_sfunc,
    FINALFUNC = pg_catalog.coord_binary_combine_agg_ffunc,
    FINALFUNC_EXTRA
);
COMMENT ON AGGREGATE pg_catalog.coord_binary_combine_agg(oid, bytea, anyelement)
    IS 'support aggregate for implementing combining partial aggregate results from workers';
REVOKE ALL ON FUNCTION pg_catalog.coord_binary_combine_agg FROM PUBLIC;
GRANT EXECUTE ON FUNCTION pg_catalog.coord_binary_combine_agg TO PUBLIC;

CREATE OR REPLACE FUNCTION pg_catalog.fix_pre_citus14_colocation_group_collation_mismatches()
RETURNS VOID AS $func$
DECLARE
    v_colocationid oid;
    v_tables_to_move_out_grouped_by_collation json;
    v_collationid oid;
    v_tables_to_move_out oid[];
    v_table_to_move_out oid;
    v_first_table oid;
BEGIN
    SET LOCAL search_path TO pg_catalog, pg_temp;
    FOR v_colocationid, v_tables_to_move_out_grouped_by_collation
    IN
       WITH colocation_groups_and_tables_with_collation_mismatches AS (
           SELECT pdc.colocationid, pa.attcollation as distkeycollation, pdp.logicalrelid
             FROM pg_dist_colocation pdc
             JOIN pg_dist_partition pdp
               ON pdc.colocationid = pdp.colocationid
             JOIN pg_attribute pa
               ON pa.attrelid = pdp.logicalrelid
              AND pa.attname = column_to_column_name(pdp.logicalrelid, pdp.partkey)
                  -- column_to_column_name() returns NULL if partkey is NULL, so we're already
                  -- implicitly ignoring the tables that don't have a distribution column, such
                  -- as reference tables, but let's still explicitly discard such tables below.
            WHERE pdp.partkey IS NOT NULL
                  -- ignore the table if its distribution column collation matches the collation saved for the colocation group
              AND pdc.distributioncolumncollation != pa.attcollation
       )
       SELECT
           colocationid,
           json_object_agg(distkeycollation, rels) AS tables_to_move_out_grouped_by_collation
       FROM (
               SELECT colocationid,
                      distkeycollation,
                      array_agg(logicalrelid::oid) AS rels
                 FROM colocation_groups_and_tables_with_collation_mismatches
             GROUP BY colocationid, distkeycollation
       ) q
       GROUP BY colocationid
    LOOP
        RAISE DEBUG 'Processing colocation group with id %', v_colocationid;
        FOR v_collationid, v_tables_to_move_out
        IN
              SELECT key::oid AS collationid,
                     array_agg(elem::oid) AS tables_to_move_out
                FROM json_each(v_tables_to_move_out_grouped_by_collation) AS e(key, value),
             LATERAL json_array_elements_text(e.value) AS elem
            GROUP BY key
        LOOP
            RAISE DEBUG 'Moving out tables with collation id % from colocation group %', v_collationid, v_colocationid;
            v_first_table := NULL;
            FOR v_table_to_move_out IN SELECT unnest(v_tables_to_move_out)
            LOOP
                IF v_first_table IS NULL then
                    -- Move the first table out to start a new colocation group.
                    --
                    -- Could check if there is an appropriate colocation group to move to instead of 'none',
                    -- but this won't be super easy. Plus, even if we had such a colocation group, the user
                    -- was anyways okay with having this in a different colocation group in the first place.
                    RAISE DEBUG 'Moving out table with oid % to a new colocation group', v_table_to_move_out;
                    PERFORM update_distributed_table_colocation(v_table_to_move_out, colocate_with => 'none');
                    -- save the first table to colocate the rest of the tables with it
                    v_first_table := v_table_to_move_out;
                ELSE
                    -- Move the rest of the tables to colocate with the first table.
                    RAISE DEBUG 'Moving out table with oid % to colocate with table with oid %', v_table_to_move_out, v_first_table;
                    PERFORM update_distributed_table_colocation(v_table_to_move_out, colocate_with => v_first_table::regclass::text);
                END IF;
            END LOOP;
        END LOOP;
    END LOOP;
END;
$func$
LANGUAGE plpgsql
SET search_path = pg_catalog, pg_temp;
COMMENT ON FUNCTION pg_catalog.fix_pre_citus14_colocation_group_collation_mismatches()
  IS 'Fix distributed tables whose colocation group collations do not match their distribution columns by moving them to new colocation groups';
CREATE OR REPLACE PROCEDURE pg_catalog.citus_finish_citus_upgrade()
    LANGUAGE plpgsql
    SET search_path = pg_catalog, pg_temp
    AS $cppu$
DECLARE
    current_version_string text;
    last_upgrade_version_string text;
    last_upgrade_major_version int;
    last_upgrade_minor_version int;
    last_upgrade_sqlpatch_version int;
    performed_upgrade bool := false;
BEGIN
 SELECT extversion INTO current_version_string
 FROM pg_extension WHERE extname = 'citus';
 -- assume some arbitrarily old version when no last upgrade version is defined
 SELECT coalesce(metadata->>'last_upgrade_version', '8.0-1') INTO last_upgrade_version_string
 FROM pg_dist_node_metadata;
 SELECT r[1], r[2], r[3]
 FROM regexp_matches(last_upgrade_version_string,'([0-9]+)\.([0-9]+)-([0-9]+)','') r
 INTO last_upgrade_major_version, last_upgrade_minor_version, last_upgrade_sqlpatch_version;
 IF last_upgrade_major_version IS NULL OR last_upgrade_minor_version IS NULL OR last_upgrade_sqlpatch_version IS NULL THEN
  -- version string is not valid, use an arbitrarily old version number
  last_upgrade_major_version := 8;
  last_upgrade_minor_version := 0;
  last_upgrade_sqlpatch_version := 1;
 END IF;
 IF last_upgrade_major_version < 11 THEN
  PERFORM citus_finalize_upgrade_to_citus11();
  performed_upgrade := true;
 END IF;
 IF last_upgrade_major_version < 14 THEN
  PERFORM fix_pre_citus14_colocation_group_collation_mismatches();
  performed_upgrade := true;
 END IF;
 -- add new upgrade steps here
 IF NOT performed_upgrade THEN
  RAISE NOTICE 'already at the latest distributed schema version (%)', last_upgrade_version_string;
  RETURN;
 END IF;
 UPDATE pg_dist_node_metadata
 SET metadata = jsonb_set(metadata, array['last_upgrade_version'], to_jsonb(current_version_string));
END;
$cppu$;
COMMENT ON PROCEDURE pg_catalog.citus_finish_citus_upgrade()
    IS 'after upgrading Citus on all nodes call this function to upgrade the distributed schema';

