From 9a2227c70da7a77e83a7103edab46199a5237762 Mon Sep 17 00:00:00 2001
From: Nils Dijk <nils@citusdata.com>
Date: Wed, 30 Jun 2021 16:27:01 +0200
Subject: [PATCH] Bump use of new sql function

---
 src/backend/distributed/citus.control          |  2 +-
 .../udfs/citus_finish_pg_upgrade/latest.sql    | 18 +++++-------------
 .../udfs/citus_prepare_pg_upgrade/latest.sql   | 10 ++++++++--
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/src/backend/distributed/citus.control b/src/backend/distributed/citus.control
index e3e408afb..631ca6850 100644
--- a/src/backend/distributed/citus.control
+++ b/src/backend/distributed/citus.control
@@ -1,6 +1,6 @@
 # Citus extension
 comment = 'Citus distributed database'
-default_version = '10.0-3'
+default_version = '10.0-4'
 module_pathname = '$libdir/citus'
 relocatable = false
 schema = pg_catalog
diff --git a/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql b/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql
index c329efea4..79153e9ae 100644
--- a/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql
+++ b/src/backend/distributed/sql/udfs/citus_finish_pg_upgrade/latest.sql
@@ -86,17 +86,7 @@ BEGIN
     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
-    )
+    TRUNCATE citus.pg_dist_object;
     INSERT INTO citus.pg_dist_object (classid, objid, objsubid, distribution_argument_index, colocationid)
     SELECT
         address.classid,
@@ -105,8 +95,10 @@ BEGIN
         naming.distribution_argument_index,
         naming.colocationid
     FROM
-        old_records naming,
-        pg_get_object_address(naming.type, naming.object_names, naming.object_args) address;
+        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;
 
     PERFORM citus_internal.columnar_ensure_objects_exist();
 END;
diff --git a/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql b/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql
index fc5dab6b1..89e3ad3c6 100644
--- a/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql
+++ b/src/backend/distributed/sql/udfs/citus_prepare_pg_upgrade/latest.sql
@@ -44,8 +44,14 @@ BEGIN
     FROM pg_catalog.pg_dist_rebalance_strategy;
 
     -- 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));
+    CREATE TABLE public.pg_dist_object AS SELECT
+       address.type,
+       address.object_names,
+       address.object_args,
+       objects.distribution_argument_index,
+       objects.colocationid
+    FROM citus.pg_dist_object objects,
+         pg_catalog.pg_identify_object_as_address(objects.classid, objects.objid, objects.objsubid) address;
 END;
 $cppu$;