Introduce a citus_finish_citus_upgrade() function

src/backend/distributed/sql/citus--11.0-1--11.0-2.sql

@@ -5,3 +5,4 @@
 #include "udfs/run_command_on_coordinator/11.0-2.sql"
 #include "udfs/start_metadata_sync_to_all_nodes/11.0-2.sql"
 #include "udfs/citus_finalize_upgrade_to_citus11/11.0-2.sql"
+#include "udfs/citus_finish_citus_upgrade/11.0-2.sql"

src/backend/distributed/sql/citus--8.0-1.sql

@@ -1461,8 +1461,17 @@ CREATE FUNCTION pg_catalog.citus_server_id()
 COMMENT ON FUNCTION citus_server_id()
     IS 'generates a random UUID to be used as server identifier';
 
+-- Insert the latest extension version into pg_dist_node_metadata
+-- for new installations.
+--
+-- While users could technically upgrade to an intermediate version
+-- everything in Citus fails until it is upgraded to the latest version,
+-- so it seems safe to use the latest.
 INSERT INTO pg_dist_node_metadata
-    VALUES (jsonb_build_object('server_id', citus_server_id()::text));
+SELECT jsonb_build_object('server_id', citus_server_id()::text,
+                          'last_upgrade_version', default_version)
+FROM pg_available_extensions
+WHERE name = 'citus';
 
 -- rebalancer functions
 CREATE TYPE citus.shard_transfer_mode AS ENUM (

src/backend/distributed/sql/downgrades/citus--11.0-2--11.0-1.sql

@@ -16,3 +16,4 @@ DROP FUNCTION pg_catalog.run_command_on_coordinator(text,boolean);
 
 DROP FUNCTION pg_catalog.start_metadata_sync_to_all_nodes();
 DROP FUNCTION pg_catalog.citus_finalize_upgrade_to_citus11(boolean);
+DROP PROCEDURE pg_catalog.citus_finish_citus_upgrade();

src/backend/distributed/sql/udfs/citus_finish_citus_upgrade/11.0-2.sql

@@ -0,0 +1,49 @@
+CREATE OR REPLACE PROCEDURE pg_catalog.citus_finish_citus_upgrade()
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    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;
+
+	-- 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';

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

@@ -0,0 +1,49 @@
+CREATE OR REPLACE PROCEDURE pg_catalog.citus_finish_citus_upgrade()
+    LANGUAGE plpgsql
+    SET search_path = pg_catalog
+    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;
+
+	-- 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';

src/test/regress/expected/multi_extension.out

@@ -1036,10 +1036,11 @@ ALTER EXTENSION citus UPDATE TO '11.0-2';
 SELECT * FROM multi_extension.print_extension_changes();
  previous_object |                         current_object
 ---------------------------------------------------------------------
+                 | function citus_finish_citus_upgrade()
                  | function citus_is_coordinator() boolean
                  | function run_command_on_coordinator(text,boolean) SETOF record
                  | function start_metadata_sync_to_all_nodes() boolean
-(3 rows)
+(4 rows)
 
 -- Test downgrade script (result should be empty)
 ALTER EXTENSION citus UPDATE TO '11.0-1';

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -60,6 +60,7 @@ ORDER BY 1;
  function citus_executor_name(integer)
  function citus_extradata_container(internal)
  function citus_finalize_upgrade_to_citus11(boolean)
+ function citus_finish_citus_upgrade()
  function citus_finish_pg_upgrade()
  function citus_get_active_worker_nodes()
  function citus_internal.columnar_ensure_am_depends_catalog()
@@ -285,5 +286,5 @@ ORDER BY 1;
  view columnar.stripe
  view pg_dist_shard_placement
  view time_partitions
-(269 rows)
+(270 rows)