Change SQL migration build process for easier reviews (#2951)

@thanodnl told me it was a bit of a problem that it's impossible to see
the history of a UDF in git. The only way to do so is by reading all the
sql migration files from new to old. Another problem is that it's also
hard to review the changed UDF during code review, because to find out
what changed you have to do the same. I thought of a IMHO better (but
not perfect) way to handle this.

We keep the definition of a UDF in sql/udfs/{name_of_udf}/latest.sql.
That file we change whenever we need to make a change to the the UDF. On
top of that you also make a snapshot of the file in
sql/udfs/{name_of_udf}/{migration-version}.sql (e.g. 9.0-1.sql) by
copying the contents. This way you can easily view what the actual
changes were by looking at the latest.sql file.

There's still the question on how to use these files then. Sadly
postgres doesn't allow inclusion of other sql files in the migration sql
file (it does in psql using \i). So instead I used the C preprocessor+
make to compile a sql/xxx.sql to a build/sql/xxx.sql file. This final
build/sql/xxx.sql file has every occurence of #include "somefile.sql" in
sql/xxx.sql replaced by the contents of somefile.sql.

.circleci/config.yml

@@ -16,6 +16,12 @@ jobs:
     steps:
       - checkout
       - {run: {name: 'Check Style', command: citus_indent --check}}
+  check-sql-snapshots:
+    docker:
+      - {image: 'citusdata/extbuilder:werror'}
+    steps:
+      - checkout
+      - {run: {name: 'Check Snapshots', command: "./configure && make check-sql-snapshots -C src/backend/distributed"}}
   test-10_check-multi:
     docker:
       - {image: 'citusdata/exttester-10:latest'}
@@ -147,6 +153,7 @@ workflows:
     jobs:
       - build
       - check-style
+      - check-sql-snapshots
 
       - {test-10_check-multi: {requires: [build]}}
       - {test-10_check-tt-van-mx: {requires: [build]}}

.gitattributes

@@ -30,3 +30,6 @@ src/backend/distributed/utils/ruleutils_10.c -citus-style
 src/backend/distributed/utils/ruleutils_11.c -citus-style
 src/backend/distributed/utils/ruleutils_12.c -citus-style
 src/include/distributed/citus_nodes.h -citus-style
+
+# Hide diff on github by default for copied udfs
+src/backend/distributed/sql/udfs/*/[123456789]*.sql linguist-generated=true

CONTRIBUTING.md

@@ -136,3 +136,39 @@ Once you've done that, you can run the `citus_indent` command to recursively che
 correct the style of any source files in the current directory. You can also run `make
 reindent` from within the Citus repo to correct the style of all source files in the
 repository.
+
+### Making SQL changes
+
+Sometimes you need to make change to the SQL that the citus extension runs upon
+creations. The way this is done is by changing the last file in
+`src/backend/distributed/sql`, or creating it if the last file is from a
+published release. If you needed to create a new file, also change the
+`default_version` field in `src/backend/distributed/citus.control` to match your
+new version. All the files in this directory are run in order based on
+their name. See [this page in the Postgres
+docs](https://www.postgresql.org/docs/current/extend-extensions.html) for more
+information on how Postgres runs these files.
+
+#### Changing or creating functions
+
+If you need to change any functions defined by Citus. You should check inside
+`src/backend/distributed/sql/udfs` to see if there is already a directory for
+this function, if not create one. Then change or create the file called
+`latest.sql` in that directory to match how it should create the function. This
+should be including any DROP (IF EXISTS), COMMENT and REVOKE statements for this
+function.
+
+Then copy the `latest.sql` file to `{version}.sql`, where `{version}` is the
+version for which this sql change is, e.g. `{9.0-1.sql}`. Now that you've
+created this stable snapshot of the function definition for your version you
+should use it in your actual sql file, .e.g.
+`src/backend/distributed/sql/citus--8.3-1--9.0-1.sql`. You do this by using C
+style `#include` statements like this:
+```
+#include "udfs/myudf/9.0-1.sql"
+```
+
+#### Other SQL
+
+Any other SQL you can put directly in the main sql file, e.g.
+`src/backend/distributed/sql/citus--8.3-1--9.0-1.sql`.

configure

@@ -1,6 +1,6 @@
 #! /bin/sh
 # Guess values for system-dependent variables and create Makefiles.
-# Generated by GNU Autoconf 2.69 for Citus 8.4devel.
+# Generated by GNU Autoconf 2.69 for Citus 9.0devel.
 #
 #
 # Copyright (C) 1992-1996, 1998-2012 Free Software Foundation, Inc.
@@ -579,8 +579,8 @@ MAKEFLAGS=
 # Identity of this package.
 PACKAGE_NAME='Citus'
 PACKAGE_TARNAME='citus'
-PACKAGE_VERSION='8.4devel'
+PACKAGE_VERSION='9.0devel'
-PACKAGE_STRING='Citus 8.4devel'
+PACKAGE_STRING='Citus 9.0devel'
 PACKAGE_BUGREPORT=''
 PACKAGE_URL=''
 
@@ -662,6 +662,7 @@ infodir
 docdir
 oldincludedir
 includedir
+runstatedir
 localstatedir
 sharedstatedir
 sysconfdir
@@ -738,6 +739,7 @@ datadir='${datarootdir}'
 sysconfdir='${prefix}/etc'
 sharedstatedir='${prefix}/com'
 localstatedir='${prefix}/var'
+runstatedir='${localstatedir}/run'
 includedir='${prefix}/include'
 oldincludedir='/usr/include'
 docdir='${datarootdir}/doc/${PACKAGE_TARNAME}'
@@ -990,6 +992,15 @@ do
   | -silent | --silent | --silen | --sile | --sil)
     silent=yes ;;
 
+  -runstatedir | --runstatedir | --runstatedi | --runstated \
+  | --runstate | --runstat | --runsta | --runst | --runs \
+  | --run | --ru | --r)
+    ac_prev=runstatedir ;;
+  -runstatedir=* | --runstatedir=* | --runstatedi=* | --runstated=* \
+  | --runstate=* | --runstat=* | --runsta=* | --runst=* | --runs=* \
+  | --run=* | --ru=* | --r=*)
+    runstatedir=$ac_optarg ;;
+
   -sbindir | --sbindir | --sbindi | --sbind | --sbin | --sbi | --sb)
     ac_prev=sbindir ;;
   -sbindir=* | --sbindir=* | --sbindi=* | --sbind=* | --sbin=* \
@@ -1127,7 +1138,7 @@ fi
 for ac_var in	exec_prefix prefix bindir sbindir libexecdir datarootdir \
 		datadir sysconfdir sharedstatedir localstatedir includedir \
 		oldincludedir docdir infodir htmldir dvidir pdfdir psdir \
-		libdir localedir mandir
+		libdir localedir mandir runstatedir
 do
   eval ac_val=\$$ac_var
   # Remove trailing slashes.
@@ -1240,7 +1251,7 @@ if test "$ac_init_help" = "long"; then
   # Omit some internal or obsolete options to make the list less imposing.
   # This message is too long to be a string in the A/UX 3.1 sh.
   cat <<_ACEOF
-\`configure' configures Citus 8.4devel to adapt to many kinds of systems.
+\`configure' configures Citus 9.0devel to adapt to many kinds of systems.
 
 Usage: $0 [OPTION]... [VAR=VALUE]...
 
@@ -1280,6 +1291,7 @@ Fine tuning of the installation directories:
   --sysconfdir=DIR        read-only single-machine data [PREFIX/etc]
   --sharedstatedir=DIR    modifiable architecture-independent data [PREFIX/com]
   --localstatedir=DIR     modifiable single-machine data [PREFIX/var]
+  --runstatedir=DIR       modifiable per-process data [LOCALSTATEDIR/run]
   --libdir=DIR            object code libraries [EPREFIX/lib]
   --includedir=DIR        C header files [PREFIX/include]
   --oldincludedir=DIR     C header files for non-gcc [/usr/include]
@@ -1301,7 +1313,7 @@ fi
 
 if test -n "$ac_init_help"; then
   case $ac_init_help in
-     short | recursive ) echo "Configuration of Citus 8.4devel:";;
+     short | recursive ) echo "Configuration of Citus 9.0devel:";;
    esac
   cat <<\_ACEOF
 
@@ -1401,7 +1413,7 @@ fi
 test -n "$ac_init_help" && exit $ac_status
 if $ac_init_version; then
   cat <<\_ACEOF
-Citus configure 8.4devel
+Citus configure 9.0devel
 generated by GNU Autoconf 2.69
 
 Copyright (C) 2012 Free Software Foundation, Inc.
@@ -1884,7 +1896,7 @@ cat >config.log <<_ACEOF
 This file contains any messages produced by compilers while
 running configure, to aid debugging if configure makes a mistake.
 
-It was created by Citus $as_me 8.4devel, which was
+It was created by Citus $as_me 9.0devel, which was
 generated by GNU Autoconf 2.69.  Invocation command line was
 
   $ $0 $@
@@ -4776,7 +4788,7 @@ cat >>$CONFIG_STATUS <<\_ACEOF || ac_write_fail=1
 # report actual input values of CONFIG_FILES etc. instead of their
 # values after options handling.
 ac_log="
-This file was extended by Citus $as_me 8.4devel, which was
+This file was extended by Citus $as_me 9.0devel, which was
 generated by GNU Autoconf 2.69.  Invocation command line was
 
   CONFIG_FILES    = $CONFIG_FILES
@@ -4838,7 +4850,7 @@ _ACEOF
 cat >>$CONFIG_STATUS <<_ACEOF || ac_write_fail=1
 ac_cs_config="`$as_echo "$ac_configure_args" | sed 's/^ //; s/[\\""\`\$]/\\\\&/g'`"
 ac_cs_version="\\
-Citus config.status 8.4devel
+Citus config.status 9.0devel
 configured by $0, generated by GNU Autoconf 2.69,
   with options \\"\$ac_cs_config\\"
 

configure.in

@@ -5,7 +5,7 @@
 # everyone needing autoconf installed, the resulting files are checked
 # into the SCM.
 
-AC_INIT([Citus], [8.4devel])
+AC_INIT([Citus], [9.0devel])
 AC_COPYRIGHT([Copyright (c) 2012-2017, Citus Data, Inc.])
 
 # we'll need sed and awk for some of the version commands

src/backend/distributed/.gitignore

@@ -8,3 +8,4 @@
 /regression.out
 /results/
 /tmp_check*
+/build/

src/backend/distributed/Makefile

@@ -6,8 +6,10 @@ citus_top_builddir = ../../..
 MODULE_big = citus
 EXTENSION = citus
 
-# All citus--*.sql files in the source directory
+template_sql_files = $(patsubst $(citus_abs_srcdir)/%,%,$(wildcard $(citus_abs_srcdir)/sql/*.sql))
-DATA = $(patsubst $(citus_abs_srcdir)/%.sql,%.sql,$(wildcard $(citus_abs_srcdir)/$(EXTENSION)--*.sql))
+generated_sql_files = $(patsubst %,$(citus_abs_srcdir)/build/%,$(template_sql_files))
+# All citus--*.sql files that are used to upgrade between versions
+DATA_built = $(generated_sql_files)
 
 # directories with source files
 SUBDIRS = . commands connection ddl deparser executor master metadata planner progress relay test transaction utils worker
@@ -29,3 +31,27 @@ include $(citus_top_builddir)/Makefile.global
 SHLIB_LINK += $(filter -lssl -lcrypto -lssleay32 -leay32, $(LIBS))
 
 override CPPFLAGS += -I$(libpq_srcdir)
+
+SQL_DEPDIR=.deps/sql
+SQL_BUILDDIR=build/sql
+
+$(SQL_DEPDIR) $(SQL_BUILDDIR):
+	mkdir -p $(citus_abs_srcdir)/$@
+
+$(generated_sql_files): $(citus_abs_srcdir)/build/%: % $(SQL_DEPDIR) $(SQL_BUILDDIR)
+	cd $(citus_abs_srcdir) && cpp -undef -w -P -MMD -MP -MF $(SQL_DEPDIR)/$(*F).Po -MT $@ $< > $@
+
+SQL_Po_files := $(wildcard $(SQL_DEPDIR)/*.Po)
+ifneq (,$(SQL_Po_files))
+include $(SQL_Po_files)
+endif
+
+.PHONY: check-sql-snapshots
+
+check-sql-snapshots:
+	bash -c '\
+	set -eu -o pipefail; \
+	for f in sql/udfs/*; do \
+		latest_snapshot=$$(ls $$f | { grep -v latest.sql || true; } | sort -V | tail -n 1); \
+		diff -u $$f/latest.sql $$f/$$latest_snapshot; \
+	done'

src/backend/distributed/citus--8.3-1--8.4-1.sql

@@ -1,372 +0,0 @@
-/* citus--8.3-1--8.4-1 */
-
-/* bump version to 8.4-1 */
-CREATE SCHEMA IF NOT EXISTS citus_internal;
-
--- move citus internal functions to citus_internal to make space in the citus schema for
--- our public interface
-ALTER FUNCTION citus.find_groupid_for_node SET SCHEMA citus_internal;
-ALTER FUNCTION citus.pg_dist_node_trigger_func SET SCHEMA citus_internal;
-ALTER FUNCTION citus.pg_dist_shard_placement_trigger_func SET SCHEMA citus_internal;
-ALTER FUNCTION citus.refresh_isolation_tester_prepared_statement SET SCHEMA citus_internal;
-ALTER FUNCTION citus.replace_isolation_tester_func SET SCHEMA citus_internal;
-ALTER FUNCTION citus.restore_isolation_tester_func SET SCHEMA citus_internal;
-
-CREATE OR REPLACE FUNCTION citus_internal.pg_dist_shard_placement_trigger_func()
-RETURNS TRIGGER AS $$
-  BEGIN
-    IF (TG_OP = 'DELETE') THEN
-      DELETE FROM pg_dist_placement WHERE placementid = OLD.placementid;
-      RETURN OLD;
-    ELSIF (TG_OP = 'UPDATE') THEN
-      UPDATE pg_dist_placement
-        SET shardid = NEW.shardid, shardstate = NEW.shardstate,
-            shardlength = NEW.shardlength, placementid = NEW.placementid,
-            groupid = citus_internal.find_groupid_for_node(NEW.nodename, NEW.nodeport)
-        WHERE placementid = OLD.placementid;
-      RETURN NEW;
-    ELSIF (TG_OP = 'INSERT') THEN
-      INSERT INTO pg_dist_placement
-        (placementid, shardid, shardstate, shardlength, groupid)
-      VALUES (NEW.placementid, NEW.shardid, NEW.shardstate, NEW.shardlength,
-        citus_internal.find_groupid_for_node(NEW.nodename, NEW.nodeport));
-      RETURN NEW;
-    END IF;
-  END;
-$$ LANGUAGE plpgsql;
-
-CREATE OR REPLACE FUNCTION pg_catalog.master_unmark_object_distributed(classid oid, objid oid, objsubid int)
-    RETURNS void
-    LANGUAGE C STRICT
-    AS 'MODULE_PATHNAME', $$master_unmark_object_distributed$$;
-COMMENT ON FUNCTION pg_catalog.master_unmark_object_distributed(classid oid, objid oid, objsubid int)
-    IS 'remove an object address from citus.pg_dist_object once the object has been deleted';
-
-CREATE OR REPLACE FUNCTION pg_catalog.worker_create_or_replace_object(statement text)
-  RETURNS bool
-  LANGUAGE C STRICT
-  AS 'MODULE_PATHNAME', $$worker_create_or_replace_object$$;
-COMMENT ON FUNCTION pg_catalog.worker_create_or_replace_object(statement text)
-    IS 'takes a sql CREATE statement, before executing the create it will check if an object with that name already exists and safely replaces that named object with the new object';
-
-CREATE TABLE citus.pg_dist_object (
-    classid oid NOT NULL,
-    objid oid NOT NULL,
-    objsubid integer NOT NULL,
-
-    -- fields used for upgrades
-    type text DEFAULT NULL,
-    object_names text[] DEFAULT NULL,
-    object_args text[] DEFAULT NULL,
-
-    CONSTRAINT pg_dist_object_pkey PRIMARY KEY (classid, objid, objsubid)
-);
-
-CREATE OR REPLACE FUNCTION pg_catalog.citus_drop_trigger()
-    RETURNS event_trigger
-    LANGUAGE plpgsql
-    SET search_path = pg_catalog
-    AS $cdbdt$
-DECLARE
-    v_obj record;
-    sequence_names text[] := '{}';
-    table_colocation_id integer;
-    propagate_drop boolean := false;
-BEGIN
-    -- collect set of dropped sequences to drop on workers later
-    SELECT array_agg(object_identity) INTO sequence_names
-    FROM pg_event_trigger_dropped_objects()
-    WHERE object_type = 'sequence';
-
-    FOR v_obj IN SELECT * FROM pg_event_trigger_dropped_objects()
-                 WHERE object_type IN ('table', 'foreign table')
-    LOOP
-        -- first drop the table and metadata on the workers
-        -- then drop all the shards on the workers
-        -- finally remove the pg_dist_partition entry on the coordinator
-        PERFORM master_remove_distributed_table_metadata_from_workers(v_obj.objid, v_obj.schema_name, v_obj.object_name);
-        PERFORM master_drop_all_shards(v_obj.objid, v_obj.schema_name, v_obj.object_name);
-        PERFORM master_remove_partition_metadata(v_obj.objid, v_obj.schema_name, v_obj.object_name);
-    END LOOP;
-
-    IF cardinality(sequence_names) > 0 THEN
-        PERFORM master_drop_sequences(sequence_names);
-    END IF;
-
-    -- remove entries from citus.pg_dist_object for all dropped root (objsubid = 0) objects
-    FOR v_obj IN SELECT * FROM pg_event_trigger_dropped_objects()
-    LOOP
-        PERFORM master_unmark_object_distributed(v_obj.classid, v_obj.objid, v_obj.objsubid);
-    END LOOP;
-END;
-$cdbdt$;
-COMMENT ON FUNCTION pg_catalog.citus_drop_trigger()
-    IS 'perform checks and actions at the end of DROP actions';
-
-
-CREATE OR REPLACE FUNCTION pg_catalog.citus_prepare_pg_upgrade()
-    RETURNS void
-    LANGUAGE plpgsql
-    SET search_path = pg_catalog
-    AS $cppu$
-BEGIN
-    --
-    -- 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;
-    -- 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;
-
-    -- store upgrade stable identifiers on pg_dist_object catalog
-    UPDATE citus.pg_dist_object
-       SET (type, object_names, object_args) = (SELECT * FROM pg_identify_object_as_address(classid, objid, objsubid));
-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
-    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;
-
-    --
-    -- 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;
-
-    --
-    -- 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
-    WITH old_records AS (
-        DELETE FROM
-            citus.pg_dist_object
-        RETURNING
-            type,
-            object_names,
-            object_args
-    )
-    INSERT INTO citus.pg_dist_object (classid, objid, objsubid)
-    SELECT
-        address.classid,
-        address.objid,
-        address.objsubid
-    FROM
-        old_records naming,
-        pg_get_object_address(naming.type, naming.object_names, naming.object_args) address;
-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';
-
-/*
- * We truncate pg_dist_node during metadata syncing, but we do not want
- * this to cascade to pg_dist_poolinfo, which is generally maintained
- * by the operator.
- */
-ALTER TABLE pg_dist_poolinfo DROP CONSTRAINT pg_dist_poolinfo_nodeid_fkey;
-
-SET search_path = 'pg_catalog';
-
-DROP EXTENSION IF EXISTS shard_rebalancer;
-
--- get_rebalance_table_shards_plan shows the actual events that will be performed
--- if a rebalance operation will be performed with the same arguments, which allows users
--- to understand the impact of the change overall availability of the application and
--- network trafic.
---
-CREATE OR REPLACE FUNCTION get_rebalance_table_shards_plan(relation regclass,
-                                                           threshold float4 default 0.1,
-                                                           max_shard_moves int default 1000000,
-                                                           excluded_shard_list bigint[] default '{}')
-    RETURNS TABLE (table_name regclass,
-                   shardid bigint,
-                   shard_size bigint,
-                   sourcename text,
-                   sourceport int,
-                   targetname text,
-                   targetport int)
-    AS 'MODULE_PATHNAME'
-    LANGUAGE C STRICT VOLATILE;
-COMMENT ON FUNCTION get_rebalance_table_shards_plan(regclass, float4, int, bigint[])
-IS 'returns the list of shard placement moves to be done on a rebalance operation';
-
--- get_rebalance_progress returns the list of shard placement move operations along with
--- their progressions for ongoing rebalance operations.
---
-CREATE OR REPLACE FUNCTION get_rebalance_progress()
-  RETURNS TABLE(sessionid integer,
-                table_name regclass,
-                shardid bigint,
-                shard_size bigint,
-                sourcename text,
-                sourceport int,
-                targetname text,
-                targetport int,
-                progress bigint)
-  AS 'MODULE_PATHNAME'
-  LANGUAGE C STRICT;
-COMMENT ON FUNCTION get_rebalance_progress()
-    IS 'provides progress information about the ongoing rebalance operations';
-
-
--- replicate_table_shards uses the shard rebalancer's C UDF functions to replicate
--- under-replicated shards of the given table.
---
-CREATE FUNCTION replicate_table_shards(relation regclass,
-                                       shard_replication_factor int default current_setting('citus.shard_replication_factor')::int,
-                                       max_shard_copies int default 1000000,
-                                       excluded_shard_list bigint[] default '{}',
-                                       shard_transfer_mode citus.shard_transfer_mode default 'auto')
-  RETURNS VOID
-  AS 'MODULE_PATHNAME'
-  LANGUAGE C STRICT;
-COMMENT ON FUNCTION replicate_table_shards(regclass, int, int, bigint[], citus.shard_transfer_mode)
-    IS 'replicates under replicated shards of the the given table';
-
--- rebalance_table_shards uses the shard rebalancer's C UDF functions to rebalance
--- shards of the given relation.
---
-CREATE OR REPLACE FUNCTION rebalance_table_shards(relation regclass,
-                                                  threshold float4 default 0,
-                                                  max_shard_moves int default 1000000,
-                                                  excluded_shard_list bigint[] default '{}',
-                                                  shard_transfer_mode citus.shard_transfer_mode default 'auto')
-  RETURNS VOID
-  AS 'MODULE_PATHNAME'
-  LANGUAGE C STRICT VOLATILE;
-COMMENT ON FUNCTION rebalance_table_shards(regclass, float4, int, bigint[], citus.shard_transfer_mode)
-    IS 'rebalance the shards of the given table across the worker nodes (including colocated shards of other tables)';
-
-DROP FUNCTION master_add_node(text, integer, integer, noderole, name);
-CREATE FUNCTION master_add_node(nodename text,
-                                nodeport integer,
-                                groupid integer default 0,
-                                noderole noderole default 'primary',
-                                nodecluster name default 'default')
-  RETURNS INTEGER
-  LANGUAGE C STRICT
-  AS 'MODULE_PATHNAME', $$master_add_node$$;
-COMMENT ON FUNCTION master_add_node(nodename text, nodeport integer,
-                                    groupid integer, noderole noderole, nodecluster name)
-  IS 'add node to the cluster';
-
-DROP FUNCTION master_add_inactive_node(text, integer, integer, noderole, name);
-CREATE FUNCTION master_add_inactive_node(nodename text,
-                                         nodeport integer,
-                                         groupid integer default 0,
-                                         noderole noderole default 'primary',
-                                         nodecluster name default 'default')
-  RETURNS INTEGER
-  LANGUAGE C STRICT
-  AS 'MODULE_PATHNAME',$$master_add_inactive_node$$;
-COMMENT ON FUNCTION master_add_inactive_node(nodename text,nodeport integer,
-                                             groupid integer, noderole noderole,
-                                             nodecluster name)
-  IS 'prepare node by adding it to pg_dist_node';
-
-SET search_path = 'pg_catalog';
-
-DROP FUNCTION master_activate_node(text, integer);
-CREATE FUNCTION master_activate_node(nodename text,
-                                     nodeport integer)
-    RETURNS INTEGER
-    LANGUAGE C STRICT
-    AS 'MODULE_PATHNAME',$$master_activate_node$$;
-COMMENT ON FUNCTION master_activate_node(nodename text, nodeport integer)
-    IS 'activate a node which is in the cluster';
-
-DROP FUNCTION master_add_secondary_node(text, integer, text, integer, name);
-CREATE FUNCTION master_add_secondary_node(nodename text,
-                                          nodeport integer,
-                                          primaryname text,
-                                          primaryport integer,
-                                          nodecluster name default 'default')
-  RETURNS INTEGER
-  LANGUAGE C STRICT
-  AS 'MODULE_PATHNAME', $$master_add_secondary_node$$;
-COMMENT ON FUNCTION master_add_secondary_node(nodename text, nodeport integer,
-                                              primaryname text, primaryport integer,
-                                              nodecluster name)
-  IS 'add a secondary node to the cluster';
-
-
-REVOKE ALL ON FUNCTION master_activate_node(text,int) FROM PUBLIC;
-REVOKE ALL ON FUNCTION master_add_inactive_node(text,int,int,noderole,name) FROM PUBLIC;
-REVOKE ALL ON FUNCTION master_add_node(text,int,int,noderole,name) FROM PUBLIC;
-REVOKE ALL ON FUNCTION master_add_secondary_node(text,int,text,int,name) FROM PUBLIC;
-
-RESET search_path;

src/backend/distributed/citus.control

@@ -1,6 +1,6 @@
 # Citus extension
 comment = 'Citus distributed database'
-default_version = '8.4-1'
+default_version = '9.0-1'
 module_pathname = '$libdir/citus'
 relocatable = false
 schema = pg_catalog

src/backend/distributed/sql/citus--7.0-4--7.0-5.sql

@@ -29,7 +29,7 @@ CREATE OR REPLACE VIEW pg_catalog.pg_dist_shard_placement AS
 CREATE OR REPLACE FUNCTION citus.pg_dist_node_trigger_func()
 RETURNS TRIGGER AS $$
   BEGIN
-    /* AddNodeMetadata also takes out a ShareRowExclusiveLock */
+    -- AddNodeMetadata also takes out a ShareRowExclusiveLock
     LOCK TABLE pg_dist_node IN SHARE ROW EXCLUSIVE MODE;
     IF (TG_OP = 'INSERT') THEN
       IF NEW.noderole = 'primary'

src/backend/distributed/sql/citus--8.3-1--9.0-1.sql

@@ -0,0 +1,139 @@
+/* citus--8.3-1--9.0-1 */
+
+/* bump version to 9.0-1 */
+CREATE SCHEMA IF NOT EXISTS citus_internal;
+
+-- move citus internal functions to citus_internal to make space in the citus schema for
+-- our public interface
+ALTER FUNCTION citus.find_groupid_for_node SET SCHEMA citus_internal;
+ALTER FUNCTION citus.pg_dist_node_trigger_func SET SCHEMA citus_internal;
+ALTER FUNCTION citus.pg_dist_shard_placement_trigger_func SET SCHEMA citus_internal;
+ALTER FUNCTION citus.refresh_isolation_tester_prepared_statement SET SCHEMA citus_internal;
+ALTER FUNCTION citus.replace_isolation_tester_func SET SCHEMA citus_internal;
+ALTER FUNCTION citus.restore_isolation_tester_func SET SCHEMA citus_internal;
+
+#include "udfs/pg_dist_shard_placement_trigger_func/9.0-1.sql"
+
+CREATE OR REPLACE FUNCTION pg_catalog.worker_create_or_replace_object(statement text)
+  RETURNS bool
+  LANGUAGE C STRICT
+  AS 'MODULE_PATHNAME', $$worker_create_or_replace_object$$;
+COMMENT ON FUNCTION pg_catalog.worker_create_or_replace_object(statement text)
+    IS 'takes a sql CREATE statement, before executing the create it will check if an object with that name already exists and safely replaces that named object with the new object';
+
+CREATE OR REPLACE FUNCTION pg_catalog.master_unmark_object_distributed(classid oid, objid oid, objsubid int)
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'MODULE_PATHNAME', $$master_unmark_object_distributed$$;
+COMMENT ON FUNCTION pg_catalog.master_unmark_object_distributed(classid oid, objid oid, objsubid int)
+    IS 'remove an object address from citus.pg_dist_object once the object has been deleted';
+
+CREATE TABLE citus.pg_dist_object (
+    classid oid NOT NULL,
+    objid oid NOT NULL,
+    objsubid integer NOT NULL,
+
+    -- fields used for upgrades
+    type text DEFAULT NULL,
+    object_names text[] DEFAULT NULL,
+    object_args text[] DEFAULT NULL,
+
+    CONSTRAINT pg_dist_object_pkey PRIMARY KEY (classid, objid, objsubid)
+);
+
+#include "udfs/citus_drop_trigger/9.0-1.sql"
+#include "udfs/citus_prepare_pg_upgrade/9.0-1.sql"
+#include "udfs/citus_finish_pg_upgrade/9.0-1.sql"
+
+/*
+ * We truncate pg_dist_node during metadata syncing, but we do not want
+ * this to cascade to pg_dist_poolinfo, which is generally maintained
+ * by the operator.
+ */
+ALTER TABLE pg_dist_poolinfo DROP CONSTRAINT pg_dist_poolinfo_nodeid_fkey;
+
+#include "udfs/get_rebalance_table_shards_plan/9.0-1.sql"
+#include "udfs/replicate_table_shards/9.0-1.sql"
+#include "udfs/rebalance_table_shards/9.0-1.sql"
+
+SET search_path = 'pg_catalog';
+
+DROP EXTENSION IF EXISTS shard_rebalancer;
+
+-- get_rebalance_progress returns the list of shard placement move operations along with
+-- their progressions for ongoing rebalance operations.
+--
+CREATE OR REPLACE FUNCTION get_rebalance_progress()
+  RETURNS TABLE(sessionid integer,
+                table_name regclass,
+                shardid bigint,
+                shard_size bigint,
+                sourcename text,
+                sourceport int,
+                targetname text,
+                targetport int,
+                progress bigint)
+  AS 'MODULE_PATHNAME'
+  LANGUAGE C STRICT;
+COMMENT ON FUNCTION get_rebalance_progress()
+    IS 'provides progress information about the ongoing rebalance operations';
+
+DROP FUNCTION master_add_node(text, integer, integer, noderole, name);
+CREATE FUNCTION master_add_node(nodename text,
+                                nodeport integer,
+                                groupid integer default 0,
+                                noderole noderole default 'primary',
+                                nodecluster name default 'default')
+  RETURNS INTEGER
+  LANGUAGE C STRICT
+  AS 'MODULE_PATHNAME', $$master_add_node$$;
+COMMENT ON FUNCTION master_add_node(nodename text, nodeport integer,
+                                    groupid integer, noderole noderole, nodecluster name)
+  IS 'add node to the cluster';
+
+DROP FUNCTION master_add_inactive_node(text, integer, integer, noderole, name);
+CREATE FUNCTION master_add_inactive_node(nodename text,
+                                         nodeport integer,
+                                         groupid integer default 0,
+                                         noderole noderole default 'primary',
+                                         nodecluster name default 'default')
+  RETURNS INTEGER
+  LANGUAGE C STRICT
+  AS 'MODULE_PATHNAME',$$master_add_inactive_node$$;
+COMMENT ON FUNCTION master_add_inactive_node(nodename text,nodeport integer,
+                                             groupid integer, noderole noderole,
+                                             nodecluster name)
+  IS 'prepare node by adding it to pg_dist_node';
+
+SET search_path = 'pg_catalog';
+
+DROP FUNCTION master_activate_node(text, integer);
+CREATE FUNCTION master_activate_node(nodename text,
+                                     nodeport integer)
+    RETURNS INTEGER
+    LANGUAGE C STRICT
+    AS 'MODULE_PATHNAME',$$master_activate_node$$;
+COMMENT ON FUNCTION master_activate_node(nodename text, nodeport integer)
+    IS 'activate a node which is in the cluster';
+
+DROP FUNCTION master_add_secondary_node(text, integer, text, integer, name);
+CREATE FUNCTION master_add_secondary_node(nodename text,
+                                          nodeport integer,
+                                          primaryname text,
+                                          primaryport integer,
+                                          nodecluster name default 'default')
+  RETURNS INTEGER
+  LANGUAGE C STRICT
+  AS 'MODULE_PATHNAME', $$master_add_secondary_node$$;
+COMMENT ON FUNCTION master_add_secondary_node(nodename text, nodeport integer,
+                                              primaryname text, primaryport integer,
+                                              nodecluster name)
+  IS 'add a secondary node to the cluster';
+
+
+REVOKE ALL ON FUNCTION master_activate_node(text,int) FROM PUBLIC;
+REVOKE ALL ON FUNCTION master_add_inactive_node(text,int,int,noderole,name) FROM PUBLIC;
+REVOKE ALL ON FUNCTION master_add_node(text,int,int,noderole,name) FROM PUBLIC;
+REVOKE ALL ON FUNCTION master_add_secondary_node(text,int,text,int,name) FROM PUBLIC;
+
+RESET search_path;

src/backend/distributed/sql/udfs/citus_drop_trigger/9.0-1.sql

@@ -0,0 +1,42 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_drop_trigger()
+    RETURNS event_trigger
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    AS $cdbdt$
+DECLARE
+    v_obj record;
+    sequence_names text[] := '{}';
+    table_colocation_id integer;
+    propagate_drop boolean := false;
+BEGIN
+    -- collect set of dropped sequences to drop on workers later
+    SELECT array_agg(object_identity) INTO sequence_names
+    FROM pg_event_trigger_dropped_objects()
+    WHERE object_type = 'sequence';
+
+    FOR v_obj IN SELECT * FROM pg_event_trigger_dropped_objects()
+                 WHERE object_type IN ('table', 'foreign table')
+    LOOP
+        -- first drop the table and metadata on the workers
+        -- then drop all the shards on the workers
+        -- finally remove the pg_dist_partition entry on the coordinator
+        PERFORM master_remove_distributed_table_metadata_from_workers(v_obj.objid, v_obj.schema_name, v_obj.object_name);
+        PERFORM master_drop_all_shards(v_obj.objid, v_obj.schema_name, v_obj.object_name);
+        PERFORM master_remove_partition_metadata(v_obj.objid, v_obj.schema_name, v_obj.object_name);
+    END LOOP;
+
+    IF cardinality(sequence_names) > 0 THEN
+        PERFORM master_drop_sequences(sequence_names);
+    END IF;
+
+    -- remove entries from citus.pg_dist_object for all dropped root (objsubid = 0) objects
+    FOR v_obj IN SELECT * FROM pg_event_trigger_dropped_objects()
+    LOOP
+        PERFORM master_unmark_object_distributed(v_obj.classid, v_obj.objid, v_obj.objsubid);
+    END LOOP;
+END;
+$cdbdt$;
+COMMENT ON FUNCTION pg_catalog.citus_drop_trigger()
+    IS 'perform checks and actions at the end of DROP actions';
+
+

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

@@ -0,0 +1,42 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_drop_trigger()
+    RETURNS event_trigger
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    AS $cdbdt$
+DECLARE
+    v_obj record;
+    sequence_names text[] := '{}';
+    table_colocation_id integer;
+    propagate_drop boolean := false;
+BEGIN
+    -- collect set of dropped sequences to drop on workers later
+    SELECT array_agg(object_identity) INTO sequence_names
+    FROM pg_event_trigger_dropped_objects()
+    WHERE object_type = 'sequence';
+
+    FOR v_obj IN SELECT * FROM pg_event_trigger_dropped_objects()
+                 WHERE object_type IN ('table', 'foreign table')
+    LOOP
+        -- first drop the table and metadata on the workers
+        -- then drop all the shards on the workers
+        -- finally remove the pg_dist_partition entry on the coordinator
+        PERFORM master_remove_distributed_table_metadata_from_workers(v_obj.objid, v_obj.schema_name, v_obj.object_name);
+        PERFORM master_drop_all_shards(v_obj.objid, v_obj.schema_name, v_obj.object_name);
+        PERFORM master_remove_partition_metadata(v_obj.objid, v_obj.schema_name, v_obj.object_name);
+    END LOOP;
+
+    IF cardinality(sequence_names) > 0 THEN
+        PERFORM master_drop_sequences(sequence_names);
+    END IF;
+
+    -- remove entries from citus.pg_dist_object for all dropped root (objsubid = 0) objects
+    FOR v_obj IN SELECT * FROM pg_event_trigger_dropped_objects()
+    LOOP
+        PERFORM master_unmark_object_distributed(v_obj.classid, v_obj.objid, v_obj.objsubid);
+    END LOOP;
+END;
+$cdbdt$;
+COMMENT ON FUNCTION pg_catalog.citus_drop_trigger()
+    IS 'perform checks and actions at the end of DROP actions';
+
+

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

@@ -0,0 +1,96 @@
+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;
+
+    --
+    -- 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;
+
+    --
+    -- 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
+    WITH old_records AS (
+        DELETE FROM
+            citus.pg_dist_object
+        RETURNING
+            type,
+            object_names,
+            object_args
+    )
+    INSERT INTO citus.pg_dist_object (classid, objid, objsubid)
+    SELECT
+        address.classid,
+        address.objid,
+        address.objsubid
+    FROM
+        old_records naming,
+        pg_get_object_address(naming.type, naming.object_names, naming.object_args) address;
+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

@@ -0,0 +1,96 @@
+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;
+
+    --
+    -- 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;
+
+    --
+    -- 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
+    WITH old_records AS (
+        DELETE FROM
+            citus.pg_dist_object
+        RETURNING
+            type,
+            object_names,
+            object_args
+    )
+    INSERT INTO citus.pg_dist_object (classid, objid, objsubid)
+    SELECT
+        address.classid,
+        address.objid,
+        address.objsubid
+    FROM
+        old_records naming,
+        pg_get_object_address(naming.type, naming.object_names, naming.object_args) address;
+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_prepare_pg_upgrade/9.0-1.sql

@@ -0,0 +1,29 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_prepare_pg_upgrade()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    AS $cppu$
+BEGIN
+    --
+    -- 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;
+    -- 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;
+
+    -- store upgrade stable identifiers on pg_dist_object catalog
+    UPDATE citus.pg_dist_object
+       SET (type, object_names, object_args) = (SELECT * FROM pg_identify_object_as_address(classid, objid, objsubid));
+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';

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

@@ -0,0 +1,29 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_prepare_pg_upgrade()
+    RETURNS void
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    AS $cppu$
+BEGIN
+    --
+    -- 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;
+    -- 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;
+
+    -- store upgrade stable identifiers on pg_dist_object catalog
+    UPDATE citus.pg_dist_object
+       SET (type, object_names, object_args) = (SELECT * FROM pg_identify_object_as_address(classid, objid, objsubid));
+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';

src/backend/distributed/sql/udfs/get_rebalance_table_shards_plan/9.0-1.sql

@@ -0,0 +1,21 @@
+-- get_rebalance_table_shards_plan shows the actual events that will be performed
+-- if a rebalance operation will be performed with the same arguments, which allows users
+-- to understand the impact of the change overall availability of the application and
+-- network trafic.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.get_rebalance_table_shards_plan(
+        relation regclass,
+        threshold float4 default 0.1,
+        max_shard_moves int default 1000000,
+        excluded_shard_list bigint[] default '{}')
+    RETURNS TABLE (table_name regclass,
+                   shardid bigint,
+                   shard_size bigint,
+                   sourcename text,
+                   sourceport int,
+                   targetname text,
+                   targetport int)
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C STRICT VOLATILE;
+COMMENT ON FUNCTION pg_catalog.get_rebalance_table_shards_plan(regclass, float4, int, bigint[])
+    IS 'returns the list of shard placement moves to be done on a rebalance operation';

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

@@ -0,0 +1,21 @@
+-- get_rebalance_table_shards_plan shows the actual events that will be performed
+-- if a rebalance operation will be performed with the same arguments, which allows users
+-- to understand the impact of the change overall availability of the application and
+-- network trafic.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.get_rebalance_table_shards_plan(
+        relation regclass,
+        threshold float4 default 0.1,
+        max_shard_moves int default 1000000,
+        excluded_shard_list bigint[] default '{}')
+    RETURNS TABLE (table_name regclass,
+                   shardid bigint,
+                   shard_size bigint,
+                   sourcename text,
+                   sourceport int,
+                   targetname text,
+                   targetport int)
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C STRICT VOLATILE;
+COMMENT ON FUNCTION pg_catalog.get_rebalance_table_shards_plan(regclass, float4, int, bigint[])
+    IS 'returns the list of shard placement moves to be done on a rebalance operation';

src/backend/distributed/sql/udfs/pg_dist_shard_placement_trigger_func/9.0-1.sql

@@ -0,0 +1,22 @@
+CREATE OR REPLACE FUNCTION citus_internal.pg_dist_shard_placement_trigger_func()
+RETURNS TRIGGER AS $$
+  BEGIN
+    IF (TG_OP = 'DELETE') THEN
+      DELETE FROM pg_dist_placement WHERE placementid = OLD.placementid;
+      RETURN OLD;
+    ELSIF (TG_OP = 'UPDATE') THEN
+      UPDATE pg_dist_placement
+        SET shardid = NEW.shardid, shardstate = NEW.shardstate,
+            shardlength = NEW.shardlength, placementid = NEW.placementid,
+            groupid = citus_internal.find_groupid_for_node(NEW.nodename, NEW.nodeport)
+        WHERE placementid = OLD.placementid;
+      RETURN NEW;
+    ELSIF (TG_OP = 'INSERT') THEN
+      INSERT INTO pg_dist_placement
+        (placementid, shardid, shardstate, shardlength, groupid)
+      VALUES (NEW.placementid, NEW.shardid, NEW.shardstate, NEW.shardlength,
+        citus_internal.find_groupid_for_node(NEW.nodename, NEW.nodeport));
+      RETURN NEW;
+    END IF;
+  END;
+$$ LANGUAGE plpgsql;

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

@@ -0,0 +1,22 @@
+CREATE OR REPLACE FUNCTION citus_internal.pg_dist_shard_placement_trigger_func()
+RETURNS TRIGGER AS $$
+  BEGIN
+    IF (TG_OP = 'DELETE') THEN
+      DELETE FROM pg_dist_placement WHERE placementid = OLD.placementid;
+      RETURN OLD;
+    ELSIF (TG_OP = 'UPDATE') THEN
+      UPDATE pg_dist_placement
+        SET shardid = NEW.shardid, shardstate = NEW.shardstate,
+            shardlength = NEW.shardlength, placementid = NEW.placementid,
+            groupid = citus_internal.find_groupid_for_node(NEW.nodename, NEW.nodeport)
+        WHERE placementid = OLD.placementid;
+      RETURN NEW;
+    ELSIF (TG_OP = 'INSERT') THEN
+      INSERT INTO pg_dist_placement
+        (placementid, shardid, shardstate, shardlength, groupid)
+      VALUES (NEW.placementid, NEW.shardid, NEW.shardstate, NEW.shardlength,
+        citus_internal.find_groupid_for_node(NEW.nodename, NEW.nodeport));
+      RETURN NEW;
+    END IF;
+  END;
+$$ LANGUAGE plpgsql;

src/backend/distributed/sql/udfs/rebalance_table_shards/9.0-1.sql

@@ -0,0 +1,14 @@
+-- rebalance_table_shards uses the shard rebalancer's C UDF functions to rebalance
+-- shards of the given relation.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.rebalance_table_shards(
+        relation regclass,
+        threshold float4 default 0,
+        max_shard_moves int default 1000000,
+        excluded_shard_list bigint[] default '{}',
+        shard_transfer_mode citus.shard_transfer_mode default 'auto')
+    RETURNS VOID
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C STRICT VOLATILE;
+COMMENT ON FUNCTION pg_catalog.rebalance_table_shards(regclass, float4, int, bigint[], citus.shard_transfer_mode)
+    IS 'rebalance the shards of the given table across the worker nodes (including colocated shards of other tables)';

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

@@ -0,0 +1,14 @@
+-- rebalance_table_shards uses the shard rebalancer's C UDF functions to rebalance
+-- shards of the given relation.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.rebalance_table_shards(
+        relation regclass,
+        threshold float4 default 0,
+        max_shard_moves int default 1000000,
+        excluded_shard_list bigint[] default '{}',
+        shard_transfer_mode citus.shard_transfer_mode default 'auto')
+    RETURNS VOID
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C STRICT VOLATILE;
+COMMENT ON FUNCTION pg_catalog.rebalance_table_shards(regclass, float4, int, bigint[], citus.shard_transfer_mode)
+    IS 'rebalance the shards of the given table across the worker nodes (including colocated shards of other tables)';

src/backend/distributed/sql/udfs/replicate_table_shards/9.0-1.sql

@@ -0,0 +1,14 @@
+-- replicate_table_shards uses the shard rebalancer's C UDF functions to replicate
+-- under-replicated shards of the given table.
+--
+CREATE FUNCTION pg_catalog.replicate_table_shards(
+        relation regclass,
+        shard_replication_factor int default current_setting('citus.shard_replication_factor')::int,
+        max_shard_copies int default 1000000,
+        excluded_shard_list bigint[] default '{}',
+        shard_transfer_mode citus.shard_transfer_mode default 'auto')
+    RETURNS VOID
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C STRICT;
+COMMENT ON FUNCTION pg_catalog.replicate_table_shards(regclass, int, int, bigint[], citus.shard_transfer_mode)
+    IS 'replicates under replicated shards of the the given table';

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

@@ -0,0 +1,14 @@
+-- replicate_table_shards uses the shard rebalancer's C UDF functions to replicate
+-- under-replicated shards of the given table.
+--
+CREATE FUNCTION pg_catalog.replicate_table_shards(
+        relation regclass,
+        shard_replication_factor int default current_setting('citus.shard_replication_factor')::int,
+        max_shard_copies int default 1000000,
+        excluded_shard_list bigint[] default '{}',
+        shard_transfer_mode citus.shard_transfer_mode default 'auto')
+    RETURNS VOID
+    AS 'MODULE_PATHNAME'
+    LANGUAGE C STRICT;
+COMMENT ON FUNCTION pg_catalog.replicate_table_shards(regclass, int, int, bigint[], citus.shard_transfer_mode)
+    IS 'replicates under replicated shards of the the given table';

src/test/regress/expected/multi_extension.out

@@ -107,12 +107,12 @@ ALTER EXTENSION citus UPDATE TO '8.2-2';
 ALTER EXTENSION citus UPDATE TO '8.2-3';
 ALTER EXTENSION citus UPDATE TO '8.2-4';
 ALTER EXTENSION citus UPDATE TO '8.3-1';
-ALTER EXTENSION citus UPDATE TO '8.4-1';
+ALTER EXTENSION citus UPDATE TO '9.0-1';
 -- show running version
 SHOW citus.version;
  citus.version 
 ---------------
- 8.4devel
+ 9.0devel
 (1 row)
 
 -- ensure no objects were created outside pg_catalog
@@ -134,7 +134,7 @@ RESET citus.enable_version_checks;
 DROP EXTENSION citus;
 CREATE EXTENSION citus VERSION '7.0-1';
 ERROR:  specified version incompatible with loaded Citus library
-DETAIL:  Loaded library requires 8.4, but 7.0-1 was specified.
+DETAIL:  Loaded library requires 9.0, but 7.0-1 was specified.
 HINT:  If a newer library is present, restart the database and try the command again.
 -- Test non-distributed queries work even in version mismatch
 SET citus.enable_version_checks TO 'false';
@@ -177,7 +177,7 @@ ORDER BY 1;
 -- We should not distribute table in version mistmatch
 SELECT create_distributed_table('version_mismatch_table', 'column1');
 ERROR:  loaded Citus library version differs from installed extension version
-DETAIL:  Loaded library requires 8.4, but the installed extension version is 7.1-1.
+DETAIL:  Loaded library requires 9.0, but the installed extension version is 7.1-1.
 HINT:  Run ALTER EXTENSION citus UPDATE and try again.
 -- This function will cause fail in next ALTER EXTENSION
 CREATE OR REPLACE FUNCTION pg_catalog.master_dist_authinfo_cache_invalidate()

src/test/regress/expected/multi_unsupported_worker_operations.out

@@ -269,7 +269,7 @@ DROP TABLE mx_table;
 ERROR:  operation is not allowed on this node
 HINT:  Connect to the coordinator and run it again.
 CONTEXT:  SQL statement "SELECT master_remove_distributed_table_metadata_from_workers(v_obj.objid, v_obj.schema_name, v_obj.object_name)"
-PL/pgSQL function citus_drop_trigger() line 19 at PERFORM
+PL/pgSQL function citus_drop_trigger() line 18 at PERFORM
 SELECT count(*) FROM mx_table;
  count 
 -------

src/test/regress/sql/multi_extension.sql

@@ -107,7 +107,7 @@ ALTER EXTENSION citus UPDATE TO '8.2-2';
 ALTER EXTENSION citus UPDATE TO '8.2-3';
 ALTER EXTENSION citus UPDATE TO '8.2-4';
 ALTER EXTENSION citus UPDATE TO '8.3-1';
-ALTER EXTENSION citus UPDATE TO '8.4-1';
+ALTER EXTENSION citus UPDATE TO '9.0-1';
 
 -- show running version
 SHOW citus.version;