create missing objects during upgrade path

ci/check_sql_snapshots.sh

@@ -4,7 +4,7 @@ set -euo pipefail
 # shellcheck disable=SC1091
 source ci/ci_helpers.sh
 
-for udf_dir in src/backend/distributed/sql/udfs/*; do
+for udf_dir in src/backend/distributed/sql/udfs/* src/backend/columnar/sql/udfs/*; do
     # We want to find the last snapshotted sql file, to make sure it's the same
     # as "latest.sql". This is done by:
     # 1. Getting the filenames in the UDF directory (using find instead of ls, to keep shellcheck happy)

src/backend/columnar/sql/columnar--9.5-1--10.0-1.sql

@@ -95,33 +95,13 @@ BEGIN
 -- user instead to add the missing objects
 IF substring(current_Setting('server_version'), '\d+')::int >= 12 THEN
   EXECUTE $$
-    CREATE FUNCTION cstore_tableam_handler(internal)
+#include "udfs/cstore_tableam_handler/10.0-1.sql"
-    RETURNS table_am_handler
+#include "udfs/alter_cstore_table_set/10.0-1.sql"
-    LANGUAGE C
+#include "udfs/alter_cstore_table_reset/10.0-1.sql"
-    AS 'MODULE_PATHNAME', 'cstore_tableam_handler';
-
-    CREATE ACCESS METHOD cstore_tableam
-    TYPE TABLE HANDLER cstore_tableam_handler;
-
-    CREATE FUNCTION pg_catalog.alter_cstore_table_set(
-        table_name regclass,
-        block_row_count int DEFAULT NULL,
-        stripe_row_count int DEFAULT NULL,
-        compression name DEFAULT null)
-    RETURNS void
-    LANGUAGE C
-    AS 'MODULE_PATHNAME', 'alter_cstore_table_set';
-
-    CREATE FUNCTION pg_catalog.alter_cstore_table_reset(
-        table_name regclass,
-        block_row_count bool DEFAULT false,
-        stripe_row_count bool DEFAULT false,
-        compression bool DEFAULT false)
-    RETURNS void
-    LANGUAGE C
-    AS 'MODULE_PATHNAME', 'alter_cstore_table_reset';
   $$;
 END IF;
 END$proc$;
 
+#include "udfs/cstore_ensure_objects_exist/10.0-1.sql"
+
 RESET search_path;

src/backend/columnar/sql/downgrades/columnar--10.0-1--9.5-1.sql

@@ -41,5 +41,7 @@ DROP FOREIGN DATA WRAPPER cstore_fdw;
 DROP FUNCTION cstore_fdw_validator(text[], oid);
 DROP FUNCTION cstore_fdw_handler();
 
+DROP FUNCTION citus_internal.cstore_ensure_objects_exist();
+
 RESET search_path;
 DROP SCHEMA cstore;

src/backend/columnar/sql/udfs/alter_cstore_table_reset/10.0-1.sql

@@ -0,0 +1,15 @@
+CREATE OR REPLACE FUNCTION pg_catalog.alter_cstore_table_reset(
+    table_name regclass,
+    block_row_count bool DEFAULT false,
+    stripe_row_count bool DEFAULT false,
+    compression bool DEFAULT false)
+    RETURNS void
+    LANGUAGE C
+AS 'MODULE_PATHNAME', 'alter_cstore_table_reset';
+
+COMMENT ON FUNCTION pg_catalog.alter_cstore_table_reset(
+    table_name regclass,
+    block_row_count bool,
+    stripe_row_count bool,
+    compression bool)
+IS 'reset on or more options on a cstore table to the system defaults';

src/backend/columnar/sql/udfs/alter_cstore_table_reset/latest.sql

@@ -0,0 +1,15 @@
+CREATE OR REPLACE FUNCTION pg_catalog.alter_cstore_table_reset(
+    table_name regclass,
+    block_row_count bool DEFAULT false,
+    stripe_row_count bool DEFAULT false,
+    compression bool DEFAULT false)
+    RETURNS void
+    LANGUAGE C
+AS 'MODULE_PATHNAME', 'alter_cstore_table_reset';
+
+COMMENT ON FUNCTION pg_catalog.alter_cstore_table_reset(
+    table_name regclass,
+    block_row_count bool,
+    stripe_row_count bool,
+    compression bool)
+IS 'reset on or more options on a cstore table to the system defaults';

src/backend/columnar/sql/udfs/alter_cstore_table_set/10.0-1.sql

@@ -0,0 +1,15 @@
+CREATE OR REPLACE FUNCTION pg_catalog.alter_cstore_table_set(
+    table_name regclass,
+    block_row_count int DEFAULT NULL,
+    stripe_row_count int DEFAULT NULL,
+    compression name DEFAULT null)
+    RETURNS void
+    LANGUAGE C
+AS 'MODULE_PATHNAME', 'alter_cstore_table_set';
+
+COMMENT ON FUNCTION pg_catalog.alter_cstore_table_set(
+    table_name regclass,
+    block_row_count int,
+    stripe_row_count int,
+    compression name)
+IS 'set one or more options on a cstore table, when set to NULL no change is made';

src/backend/columnar/sql/udfs/alter_cstore_table_set/latest.sql

@@ -0,0 +1,15 @@
+CREATE OR REPLACE FUNCTION pg_catalog.alter_cstore_table_set(
+    table_name regclass,
+    block_row_count int DEFAULT NULL,
+    stripe_row_count int DEFAULT NULL,
+    compression name DEFAULT null)
+    RETURNS void
+    LANGUAGE C
+AS 'MODULE_PATHNAME', 'alter_cstore_table_set';
+
+COMMENT ON FUNCTION pg_catalog.alter_cstore_table_set(
+    table_name regclass,
+    block_row_count int,
+    stripe_row_count int,
+    compression name)
+IS 'set one or more options on a cstore table, when set to NULL no change is made';

src/backend/columnar/sql/udfs/cstore_ensure_objects_exist/10.0-1.sql

@@ -0,0 +1,47 @@
+-- citus_internal.cstore_ensure_objects_exist is an internal helper function to create
+-- missing objects, anything related to the table access methods.
+-- Since the API for table access methods is only available in PG12 we can't create these
+-- objects when Citus is installed in PG11. Once citus is installed on PG11 the user can
+-- upgrade their database to PG12. Now they require the table access method objects that
+-- we couldn't create before.
+-- This internal function is called from `citus_finish_pg_upgrade` which the user is
+-- required to call after a PG major upgrade.
+CREATE OR REPLACE FUNCTION citus_internal.cstore_ensure_objects_exist()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+AS $ceoe$
+BEGIN
+
+-- when postgres is version 12 or above we need to create the tableam. If the tableam
+-- exist we assume all objects have been created.
+IF substring(current_Setting('server_version'), '\d+')::int >= 12 THEN
+IF NOT EXISTS (SELECT 1 FROM pg_am WHERE amname = 'cstore_tableam') THEN
+
+#include "../cstore_tableam_handler/10.0-1.sql"
+
+#include "../alter_cstore_table_set/10.0-1.sql"
+
+#include "../alter_cstore_table_reset/10.0-1.sql"
+
+    -- add the missing objects to the extension
+    ALTER EXTENSION citus ADD FUNCTION cstore.cstore_tableam_handler(internal);
+    ALTER EXTENSION citus ADD ACCESS METHOD cstore_tableam;
+    ALTER EXTENSION citus ADD FUNCTION pg_catalog.alter_cstore_table_set(
+        table_name regclass,
+        block_row_count int,
+        stripe_row_count int,
+        compression name);
+    ALTER EXTENSION citus ADD FUNCTION pg_catalog.alter_cstore_table_reset(
+        table_name regclass,
+        block_row_count bool,
+        stripe_row_count bool,
+        compression bool);
+
+END IF;
+END IF;
+END;
+$ceoe$;
+
+COMMENT ON FUNCTION citus_internal.cstore_ensure_objects_exist()
+    IS 'internal function to be called by pg_catalog.citus_finish_pg_upgrade responsible for creating the columnar objects';

src/backend/columnar/sql/udfs/cstore_ensure_objects_exist/latest.sql

@@ -0,0 +1,47 @@
+-- citus_internal.cstore_ensure_objects_exist is an internal helper function to create
+-- missing objects, anything related to the table access methods.
+-- Since the API for table access methods is only available in PG12 we can't create these
+-- objects when Citus is installed in PG11. Once citus is installed on PG11 the user can
+-- upgrade their database to PG12. Now they require the table access method objects that
+-- we couldn't create before.
+-- This internal function is called from `citus_finish_pg_upgrade` which the user is
+-- required to call after a PG major upgrade.
+CREATE OR REPLACE FUNCTION citus_internal.cstore_ensure_objects_exist()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+AS $ceoe$
+BEGIN
+
+-- when postgres is version 12 or above we need to create the tableam. If the tableam
+-- exist we assume all objects have been created.
+IF substring(current_Setting('server_version'), '\d+')::int >= 12 THEN
+IF NOT EXISTS (SELECT 1 FROM pg_am WHERE amname = 'cstore_tableam') THEN
+
+#include "../cstore_tableam_handler/10.0-1.sql"
+
+#include "../alter_cstore_table_set/10.0-1.sql"
+
+#include "../alter_cstore_table_reset/10.0-1.sql"
+
+    -- add the missing objects to the extension
+    ALTER EXTENSION citus ADD FUNCTION cstore.cstore_tableam_handler(internal);
+    ALTER EXTENSION citus ADD ACCESS METHOD cstore_tableam;
+    ALTER EXTENSION citus ADD FUNCTION pg_catalog.alter_cstore_table_set(
+        table_name regclass,
+        block_row_count int,
+        stripe_row_count int,
+        compression name);
+    ALTER EXTENSION citus ADD FUNCTION pg_catalog.alter_cstore_table_reset(
+        table_name regclass,
+        block_row_count bool,
+        stripe_row_count bool,
+        compression bool);
+
+END IF;
+END IF;
+END;
+$ceoe$;
+
+COMMENT ON FUNCTION citus_internal.cstore_ensure_objects_exist()
+    IS 'internal function to be called by pg_catalog.citus_finish_pg_upgrade responsible for creating the columnar objects';

src/backend/columnar/sql/udfs/cstore_tableam_handler/10.0-1.sql

@@ -0,0 +1,9 @@
+CREATE OR REPLACE FUNCTION cstore.cstore_tableam_handler(internal)
+    RETURNS table_am_handler
+    LANGUAGE C
+AS 'MODULE_PATHNAME', 'cstore_tableam_handler';
+
+COMMENT ON FUNCTION cstore.cstore_tableam_handler(internal)
+    IS 'internal function returning the handler for cstore tables';
+
+CREATE ACCESS METHOD cstore_tableam TYPE TABLE HANDLER cstore.cstore_tableam_handler;

src/backend/columnar/sql/udfs/cstore_tableam_handler/latest.sql

@@ -0,0 +1,9 @@
+CREATE OR REPLACE FUNCTION cstore.cstore_tableam_handler(internal)
+    RETURNS table_am_handler
+    LANGUAGE C
+AS 'MODULE_PATHNAME', 'cstore_tableam_handler';
+
+COMMENT ON FUNCTION cstore.cstore_tableam_handler(internal)
+    IS 'internal function returning the handler for cstore tables';
+
+CREATE ACCESS METHOD cstore_tableam TYPE TABLE HANDLER cstore.cstore_tableam_handler;

src/backend/distributed/sql/citus--9.5-1--10.0-1.sql

@@ -2,4 +2,6 @@
 
 -- bump version to 10.0-1
 
+#include "udfs/citus_finish_pg_upgrade/10.0-1.sql"
+
 #include "../../columnar/sql/columnar--9.5-1--10.0-1.sql"

src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql

@@ -1,4 +1,6 @@
 -- citus--10.0-1--9.5-1
 -- this is an empty downgrade path since citus--9.5-1--10.0-1.sql is empty for now
 
+#include "../udfs/citus_finish_pg_upgrade/9.5-1.sql"
+
 #include "../../../columnar/sql/downgrades/columnar--10.0-1--9.5-1.sql"

src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/10.0-1.sql

@@ -0,0 +1,116 @@
+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
+    --
+    -- restore citus catalog tables
+    --
+    INSERT INTO pg_catalog.pg_dist_partition SELECT * FROM public.pg_dist_partition;
+    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;
+    -- 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;
+
+    ALTER TABLE pg_catalog.pg_dist_rebalance_strategy DISABLE TRIGGER pg_dist_rebalance_strategy_enterprise_check_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
+    FROM public.pg_dist_rebalance_strategy;
+    ALTER TABLE pg_catalog.pg_dist_rebalance_strategy ENABLE TRIGGER pg_dist_rebalance_strategy_enterprise_check_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;
+
+    --
+    -- 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);
+
+    --
+    -- register triggers
+    --
+    FOR table_name IN SELECT logicalrelid FROM pg_catalog.pg_dist_partition
+    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;
+
+    -- restore pg_dist_object from the stable identifiers
+    -- DELETE/INSERT to avoid primary key violations
+    WITH old_records AS (
+        DELETE FROM
+            citus.pg_dist_object
+        RETURNING
+            type,
+            object_names,
+            object_args,
+            distribution_argument_index,
+            colocationid
+    )
+    INSERT INTO citus.pg_dist_object (classid, objid, objsubid, distribution_argument_index, colocationid)
+    SELECT
+        address.classid,
+        address.objid,
+        address.objsubid,
+        naming.distribution_argument_index,
+        naming.colocationid
+    FROM
+        old_records naming,
+        pg_get_object_address(naming.type, naming.object_names, naming.object_args) address;
+
+    PERFORM citus_internal.cstore_ensure_objects_exist();
+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';

src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql

@@ -107,6 +107,8 @@ BEGIN
     FROM
         old_records naming,
         pg_get_object_address(naming.type, naming.object_names, naming.object_args) address;
+
+    PERFORM citus_internal.cstore_ensure_objects_exist();
 END;
 $cppu$;
 

src/test/regress/expected/multi_extension.out

@@ -472,6 +472,7 @@ SELECT * FROM print_extension_changes();
                  | foreign-data wrapper cstore_fdw
                  | function alter_cstore_table_reset(regclass,boolean,boolean,boolean)
                  | function alter_cstore_table_set(regclass,integer,integer,name)
+                 | function citus_internal.cstore_ensure_objects_exist()
                  | function cstore.cstore_ddl_event_end_trigger()
                  | function cstore.cstore_fdw_handler()
                  | function cstore.cstore_fdw_validator(text[],oid)
@@ -482,7 +483,7 @@ SELECT * FROM print_extension_changes();
                  | table cstore.cstore_skipnodes
                  | table cstore.cstore_stripes
                  | view cstore.cstore_options
-(15 rows)
+(16 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -38,6 +38,7 @@ ORDER BY 1;
  function citus_executor_name(integer)
  function citus_extradata_container(internal)
  function citus_finish_pg_upgrade()
+ function citus_internal.cstore_ensure_objects_exist()
  function citus_internal.find_groupid_for_node(text,integer)
  function citus_internal.pg_dist_node_trigger_func()
  function citus_internal.pg_dist_rebalance_strategy_enterprise_check()
@@ -222,5 +223,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view cstore.cstore_options
  view pg_dist_shard_placement
-(206 rows)
+(207 rows)
 

src/test/regress/expected/upgrade_list_citus_objects_0.out

@@ -35,6 +35,7 @@ ORDER BY 1;
  function citus_executor_name(integer)
  function citus_extradata_container(internal)
  function citus_finish_pg_upgrade()
+ function citus_internal.cstore_ensure_objects_exist()
  function citus_internal.find_groupid_for_node(text,integer)
  function citus_internal.pg_dist_node_trigger_func()
  function citus_internal.pg_dist_rebalance_strategy_enterprise_check()
@@ -218,5 +219,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view cstore.cstore_options
  view pg_dist_shard_placement
-(202 rows)
+(203 rows)