From 18d8a72ca8da5916098320ae69cf15f7d9ed4735 Mon Sep 17 00:00:00 2001
From: Metin Doslu <metin@citusdata.com>
Date: Tue, 6 Dec 2016 15:40:21 +0200
Subject: [PATCH] Add mark_colocation_group_default()

This udf helps to update default colocation group.
---
 .../distributed/citus--6.1-3--6.1-4.sql       | 47 +++++++++++++++++++
 .../expected/multi_colocation_utils.out       |  4 +-
 .../regress/sql/multi_colocation_utils.sql    |  4 +-
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/src/backend/distributed/citus--6.1-3--6.1-4.sql b/src/backend/distributed/citus--6.1-3--6.1-4.sql
index 15294e3b2..89d7c741e 100644
--- a/src/backend/distributed/citus--6.1-3--6.1-4.sql
+++ b/src/backend/distributed/citus--6.1-3--6.1-4.sql
@@ -26,4 +26,51 @@ COMMENT ON FUNCTION create_distributed_table(table_name regclass,
 											 colocate_with text)
     IS 'creates a distributed table';
 
+CREATE FUNCTION mark_colocation_group_default(colocation_id integer)
+	RETURNS void
+	LANGUAGE plpgsql
+	SET search_path = pg_catalog
+	AS $$
+DECLARE
+	shard_count integer;
+	replication_factor integer;
+	distribution_column_type oid;
+BEGIN
+	-- get colocation group configuration
+	SELECT
+		shardcount,
+		replicationfactor,
+		distributioncolumntype
+	INTO
+		shard_count,
+		replication_factor,
+		distribution_column_type
+	FROM
+		pg_dist_colocation
+	WHERE
+		colocationid = colocation_id;
+
+	-- set all defaults to false
+	UPDATE
+		pg_dist_colocation
+	SET
+		defaultgroup = false
+	WHERE
+		shardcount = shard_count AND
+		replicationfactor = replication_factor AND
+		distributioncolumntype = distribution_column_type;
+
+	-- set new default colocation group
+	UPDATE
+		pg_dist_colocation
+	SET
+		defaultgroup = true
+	WHERE
+		colocationid = colocation_id;
+END;
+$$;
+
+COMMENT ON FUNCTION mark_colocation_group_default(colocation_id integer)
+    IS 'marks given colocation group as default';
+
 RESET search_path;
diff --git a/src/test/regress/expected/multi_colocation_utils.out b/src/test/regress/expected/multi_colocation_utils.out
index 933e8615d..cc5d485af 100644
--- a/src/test/regress/expected/multi_colocation_utils.out
+++ b/src/test/regress/expected/multi_colocation_utils.out
@@ -871,8 +871,8 @@ DELETE FROM pg_dist_colocation
 UPDATE pg_dist_partition SET colocationid = 0
     WHERE colocationid >= 1 AND colocationid < 1000;
 -- check metadata
-SELECT * FROM pg_dist_colocation 
-    WHERE colocationid >= 1 AND colocationid < 1000 
+SELECT * FROM pg_dist_colocation
+    WHERE colocationid >= 1 AND colocationid < 1000
     ORDER BY colocationid;
  colocationid | shardcount | replicationfactor | distributioncolumntype | defaultgroup 
 --------------+------------+-------------------+------------------------+--------------
diff --git a/src/test/regress/sql/multi_colocation_utils.sql b/src/test/regress/sql/multi_colocation_utils.sql
index 53d761b73..3f65de670 100644
--- a/src/test/regress/sql/multi_colocation_utils.sql
+++ b/src/test/regress/sql/multi_colocation_utils.sql
@@ -302,8 +302,8 @@ CREATE TABLE table2_groupF ( id int );
 SELECT create_reference_table('table2_groupF');
 
 -- check metadata
-SELECT * FROM pg_dist_colocation 
-    WHERE colocationid >= 1 AND colocationid < 1000 
+SELECT * FROM pg_dist_colocation
+    WHERE colocationid >= 1 AND colocationid < 1000
     ORDER BY colocationid;
 
 -- test mark_colocation_group_default()