From ceb6b503c01a481901d0af707f3b41b6cc9722fc Mon Sep 17 00:00:00 2001
From: Ahmet Gedemenli <afgedemenli@gmail.com>
Date: Mon, 18 Jan 2021 12:25:54 +0300
Subject: [PATCH 1/2] Remove unused UDF mark_tables_colocated

---
 src/backend/distributed/sql/citus--9.5-1--10.0-1.sql       | 1 +
 .../distributed/sql/downgrades/citus--10.0-1--9.5-1.sql    | 7 +++++++
 src/test/regress/expected/multi_drop_extension.out         | 3 ++-
 src/test/regress/expected/multi_extension.out              | 3 ++-
 src/test/regress/expected/multi_extension_0.out            | 3 ++-
 src/test/regress/expected/multi_metadata_sync.out          | 5 +++++
 src/test/regress/expected/multi_reference_table.out        | 5 +++++
 .../expected/multi_unsupported_worker_operations.out       | 5 +++++
 src/test/regress/expected/upgrade_list_citus_objects.out   | 3 +--
 src/test/regress/expected/upgrade_list_citus_objects_0.out | 3 +--
 src/test/regress/sql/multi_drop_extension.sql              | 3 ++-
 src/test/regress/sql/multi_metadata_sync.sql               | 5 +++++
 src/test/regress/sql/multi_reference_table.sql             | 6 ++++++
 .../regress/sql/multi_unsupported_worker_operations.sql    | 5 +++++
 14 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/src/backend/distributed/sql/citus--9.5-1--10.0-1.sql b/src/backend/distributed/sql/citus--9.5-1--10.0-1.sql
index be3ee5db9..03b0ff06c 100644
--- a/src/backend/distributed/sql/citus--9.5-1--10.0-1.sql
+++ b/src/backend/distributed/sql/citus--9.5-1--10.0-1.sql
@@ -63,6 +63,7 @@ RENAME TO citus_drop_all_shards;
 DROP FUNCTION pg_catalog.master_modify_multiple_shards(text);
 DROP FUNCTION pg_catalog.master_create_distributed_table(regclass, text, citus.distribution_type);
 DROP FUNCTION pg_catalog.master_create_worker_shards(text, integer, integer);
+DROP FUNCTION pg_catalog.mark_tables_colocated(regclass, regclass[]);
 #include "udfs/citus_shard_sizes/10.0-1.sql"
 #include "udfs/citus_shards/10.0-1.sql"
 
diff --git a/src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql b/src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql
index 1fbbf1f43..a96906004 100644
--- a/src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql
+++ b/src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql
@@ -58,6 +58,13 @@ DROP PROCEDURE pg_catalog.alter_old_partitions_set_access_method(regclass,timest
 DROP FUNCTION pg_catalog.citus_set_coordinator_host(text,int,noderole,name);
 DROP FUNCTION pg_catalog.worker_change_sequence_dependency(regclass, regclass, regclass);
 
+CREATE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'MODULE_PATHNAME', $$mark_tables_colocated$$;
+COMMENT ON FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	IS 'mark target distributed tables as colocated with the source table';
+
 CREATE FUNCTION pg_catalog.master_modify_multiple_shards(text)
     RETURNS integer
     LANGUAGE C STRICT
diff --git a/src/test/regress/expected/multi_drop_extension.out b/src/test/regress/expected/multi_drop_extension.out
index afa9a4ed2..cd3098ee8 100644
--- a/src/test/regress/expected/multi_drop_extension.out
+++ b/src/test/regress/expected/multi_drop_extension.out
@@ -15,7 +15,8 @@ SELECT create_distributed_table('testtableddl', 'distributecol', 'append');
 -- change this test every time the previous tests change the set of tables they leave
 -- around.
 SET client_min_messages TO 'WARNING';
-DROP FUNCTION pg_catalog.master_create_worker_shards;
+DROP FUNCTION pg_catalog.master_create_worker_shards(text, integer, integer);
+DROP FUNCTION pg_catalog.mark_tables_colocated(regclass, regclass[]);
 DROP EXTENSION citus CASCADE;
 RESET client_min_messages;
 BEGIN;
diff --git a/src/test/regress/expected/multi_extension.out b/src/test/regress/expected/multi_extension.out
index a56701ca9..fe1b3a195 100644
--- a/src/test/regress/expected/multi_extension.out
+++ b/src/test/regress/expected/multi_extension.out
@@ -447,6 +447,7 @@ SELECT * FROM print_extension_changes();
 ---------------------------------------------------------------------
  function citus_total_relation_size(regclass)                                    |
  function create_citus_local_table(regclass)                                     |
+ function mark_tables_colocated(regclass, regclass[])                            |
  function master_conninfo_cache_invalidate()                                     |
  function master_create_distributed_table(regclass,text,citus.distribution_type) |
  function master_create_worker_shards(text,integer,integer)                      |
@@ -505,7 +506,7 @@ SELECT * FROM print_extension_changes();
                                                                                  | view citus_shards
                                                                                  | view citus_tables
                                                                                  | view time_partitions
-(60 rows)
+(61 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version
diff --git a/src/test/regress/expected/multi_extension_0.out b/src/test/regress/expected/multi_extension_0.out
index 89d4382b3..5ec7f3990 100644
--- a/src/test/regress/expected/multi_extension_0.out
+++ b/src/test/regress/expected/multi_extension_0.out
@@ -447,6 +447,7 @@ SELECT * FROM print_extension_changes();
 ---------------------------------------------------------------------
  function citus_total_relation_size(regclass)                                    |
  function create_citus_local_table(regclass)                                     |
+ function mark_tables_colocated(regclass, regclass[])                            |
  function master_conninfo_cache_invalidate()                                     |
  function master_create_distributed_table(regclass,text,citus.distribution_type) |
  function master_create_worker_shards(text,integer,integer)                      |
@@ -501,7 +502,7 @@ SELECT * FROM print_extension_changes();
                                                                                  | view citus_shards
                                                                                  | view citus_tables
                                                                                  | view time_partitions
-(56 rows)
+(57 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version
diff --git a/src/test/regress/expected/multi_metadata_sync.out b/src/test/regress/expected/multi_metadata_sync.out
index f3aeede8d..c457b5a6b 100644
--- a/src/test/regress/expected/multi_metadata_sync.out
+++ b/src/test/regress/expected/multi_metadata_sync.out
@@ -810,6 +810,11 @@ SELECT "Constraint", "Definition" FROM table_fkeys WHERE relid='mx_test_schema_1
 
 -- Check that mark_tables_colocated call propagates the changes to the workers
 \c - - - :master_port
+-- this function is dropped in Citus10, added here for tests
+CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'citus', $$mark_tables_colocated$$;
 SELECT nextval('pg_catalog.pg_dist_colocationid_seq') AS last_colocation_id \gset
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 10000;
 SET citus.shard_count TO 7;
diff --git a/src/test/regress/expected/multi_reference_table.out b/src/test/regress/expected/multi_reference_table.out
index 70da4735e..83e7e8c84 100644
--- a/src/test/regress/expected/multi_reference_table.out
+++ b/src/test/regress/expected/multi_reference_table.out
@@ -1196,6 +1196,11 @@ DEBUG:  cannot perform distributed INSERT INTO ... SELECT because the partition
 DETAIL:  The target table's partition column should correspond to a partition column in the subquery.
 DEBUG:  performing repartitioned INSERT ... SELECT
 RESET client_min_messages;
+-- this function is dropped in Citus10, added here for tests
+CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'citus', $$mark_tables_colocated$$;
 -- some tests for mark_tables_colocated
 -- should error out
 SELECT mark_tables_colocated('colocated_table_test_2', ARRAY['reference_table_test']);
diff --git a/src/test/regress/expected/multi_unsupported_worker_operations.out b/src/test/regress/expected/multi_unsupported_worker_operations.out
index 5f92cfd8f..7ec8728bd 100644
--- a/src/test/regress/expected/multi_unsupported_worker_operations.out
+++ b/src/test/regress/expected/multi_unsupported_worker_operations.out
@@ -221,6 +221,11 @@ SELECT master_remove_node('localhost', 5432);
 
 \c - - - :worker_1_port
 -- mark_tables_colocated
+-- this function is dropped in Citus10, added here for tests
+CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'citus', $$mark_tables_colocated$$;
 UPDATE pg_dist_partition SET colocationid = 0 WHERE logicalrelid='mx_table_2'::regclass;
 SELECT mark_tables_colocated('mx_table', ARRAY['mx_table_2']);
 ERROR:  operation is not allowed on this node
diff --git a/src/test/regress/expected/upgrade_list_citus_objects.out b/src/test/regress/expected/upgrade_list_citus_objects.out
index 00cb067f5..db85ce4bb 100644
--- a/src/test/regress/expected/upgrade_list_citus_objects.out
+++ b/src/test/regress/expected/upgrade_list_citus_objects.out
@@ -126,7 +126,6 @@ ORDER BY 1;
  function lock_relation_if_exists(text,text)
  function lock_shard_metadata(integer,bigint[])
  function lock_shard_resources(integer,bigint[])
- function mark_tables_colocated(regclass,regclass[])
  function master_activate_node(text,integer)
  function master_add_inactive_node(text,integer,integer,noderole,name)
  function master_add_node(text,integer,integer,noderole,name)
@@ -237,5 +236,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view pg_dist_shard_placement
  view time_partitions
-(221 rows)
+(220 rows)
 
diff --git a/src/test/regress/expected/upgrade_list_citus_objects_0.out b/src/test/regress/expected/upgrade_list_citus_objects_0.out
index d9980d8bf..a7fcf6185 100644
--- a/src/test/regress/expected/upgrade_list_citus_objects_0.out
+++ b/src/test/regress/expected/upgrade_list_citus_objects_0.out
@@ -122,7 +122,6 @@ ORDER BY 1;
  function lock_relation_if_exists(text,text)
  function lock_shard_metadata(integer,bigint[])
  function lock_shard_resources(integer,bigint[])
- function mark_tables_colocated(regclass,regclass[])
  function master_activate_node(text,integer)
  function master_add_inactive_node(text,integer,integer,noderole,name)
  function master_add_node(text,integer,integer,noderole,name)
@@ -233,5 +232,5 @@ ORDER BY 1;
  view citus_worker_stat_activity
  view pg_dist_shard_placement
  view time_partitions
-(217 rows)
+(216 rows)
 
diff --git a/src/test/regress/sql/multi_drop_extension.sql b/src/test/regress/sql/multi_drop_extension.sql
index 8112ffef1..ec1bee234 100644
--- a/src/test/regress/sql/multi_drop_extension.sql
+++ b/src/test/regress/sql/multi_drop_extension.sql
@@ -15,7 +15,8 @@ SELECT create_distributed_table('testtableddl', 'distributecol', 'append');
 -- change this test every time the previous tests change the set of tables they leave
 -- around.
 SET client_min_messages TO 'WARNING';
-DROP FUNCTION pg_catalog.master_create_worker_shards;
+DROP FUNCTION pg_catalog.master_create_worker_shards(text, integer, integer);
+DROP FUNCTION pg_catalog.mark_tables_colocated(regclass, regclass[]);
 DROP EXTENSION citus CASCADE;
 RESET client_min_messages;
 
diff --git a/src/test/regress/sql/multi_metadata_sync.sql b/src/test/regress/sql/multi_metadata_sync.sql
index 425685fda..a0ff53a7c 100644
--- a/src/test/regress/sql/multi_metadata_sync.sql
+++ b/src/test/regress/sql/multi_metadata_sync.sql
@@ -356,6 +356,11 @@ SELECT "Constraint", "Definition" FROM table_fkeys WHERE relid='mx_test_schema_1
 
 -- Check that mark_tables_colocated call propagates the changes to the workers
 \c - - - :master_port
+-- this function is dropped in Citus10, added here for tests
+CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'citus', $$mark_tables_colocated$$;
 SELECT nextval('pg_catalog.pg_dist_colocationid_seq') AS last_colocation_id \gset
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 10000;
 SET citus.shard_count TO 7;
diff --git a/src/test/regress/sql/multi_reference_table.sql b/src/test/regress/sql/multi_reference_table.sql
index aec36920b..34b9900db 100644
--- a/src/test/regress/sql/multi_reference_table.sql
+++ b/src/test/regress/sql/multi_reference_table.sql
@@ -764,6 +764,12 @@ WHERE
 
 RESET client_min_messages;
 
+-- this function is dropped in Citus10, added here for tests
+CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'citus', $$mark_tables_colocated$$;
+
 -- some tests for mark_tables_colocated
 -- should error out
 SELECT mark_tables_colocated('colocated_table_test_2', ARRAY['reference_table_test']);
diff --git a/src/test/regress/sql/multi_unsupported_worker_operations.sql b/src/test/regress/sql/multi_unsupported_worker_operations.sql
index 22e8cc5dd..69677d3c1 100644
--- a/src/test/regress/sql/multi_unsupported_worker_operations.sql
+++ b/src/test/regress/sql/multi_unsupported_worker_operations.sql
@@ -143,6 +143,11 @@ SELECT master_remove_node('localhost', 5432);
 \c - - - :worker_1_port
 
 -- mark_tables_colocated
+-- this function is dropped in Citus10, added here for tests
+CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
+	RETURNS void
+	LANGUAGE C STRICT
+	AS 'citus', $$mark_tables_colocated$$;
 UPDATE pg_dist_partition SET colocationid = 0 WHERE logicalrelid='mx_table_2'::regclass;
 
 SELECT mark_tables_colocated('mx_table', ARRAY['mx_table_2']);

From 89a6fe83f7e9bdc35847a7a9674b57ca070697c7 Mon Sep 17 00:00:00 2001
From: Ahmet Gedemenli <afgedemenli@gmail.com>
Date: Tue, 19 Jan 2021 14:01:18 +0300
Subject: [PATCH 2/2] Replace to update_distributed_table_colocation for tests

---
 .../distributed/utils/colocation_utils.c      |   6 +-
 .../regress/expected/citus_local_tables.out   |  10 +-
 .../expected/multi_colocation_utils.out       | 267 +++++-------------
 .../regress/expected/multi_drop_extension.out |   1 -
 .../regress/expected/multi_metadata_sync.out  |  17 +-
 src/test/regress/expected/multi_multiuser.out |   2 +-
 .../expected/multi_reference_table.out        |  23 +-
 .../multi_unsupported_worker_operations.out   |   8 +-
 src/test/regress/sql/citus_local_tables.sql   |  10 +-
 .../regress/sql/multi_colocation_utils.sql    |  64 ++---
 src/test/regress/sql/multi_drop_extension.sql |   1 -
 src/test/regress/sql/multi_metadata_sync.sql  |  11 +-
 src/test/regress/sql/multi_multiuser.sql      |   2 +-
 .../regress/sql/multi_reference_table.sql     |  11 +-
 .../multi_unsupported_worker_operations.sql   |   8 +-
 15 files changed, 127 insertions(+), 314 deletions(-)

diff --git a/src/backend/distributed/utils/colocation_utils.c b/src/backend/distributed/utils/colocation_utils.c
index b1038fd5f..d1307c551 100644
--- a/src/backend/distributed/utils/colocation_utils.c
+++ b/src/backend/distributed/utils/colocation_utils.c
@@ -113,17 +113,17 @@ update_distributed_table_colocation(PG_FUNCTION_ARGS)
 	CheckCitusVersion(ERROR);
 	EnsureCoordinator();
 	EnsureTableOwner(targetRelationId);
-	EnsureHashDistributedTable(targetRelationId);
 
 	char *colocateWithTableName = text_to_cstring(colocateWithTableNameText);
 	if (IsColocateWithNone(colocateWithTableName))
 	{
+		EnsureHashDistributedTable(targetRelationId);
 		BreakColocation(targetRelationId);
 	}
 	else
 	{
 		Oid colocateWithTableId = ResolveRelationId(colocateWithTableNameText, false);
-		EnsureHashDistributedTable(colocateWithTableId);
+		EnsureTableOwner(colocateWithTableId);
 		MarkTablesColocated(colocateWithTableId, targetRelationId);
 	}
 	PG_RETURN_VOID();
@@ -248,6 +248,8 @@ MarkTablesColocated(Oid sourceRelationId, Oid targetRelationId)
 							   "other tables")));
 	}
 
+	EnsureHashDistributedTable(sourceRelationId);
+	EnsureHashDistributedTable(targetRelationId);
 	CheckReplicationModel(sourceRelationId, targetRelationId);
 	CheckDistributionColumnType(sourceRelationId, targetRelationId);
 
diff --git a/src/test/regress/expected/citus_local_tables.out b/src/test/regress/expected/citus_local_tables.out
index 65dc6fe1d..45beceeef 100644
--- a/src/test/regress/expected/citus_local_tables.out
+++ b/src/test/regress/expected/citus_local_tables.out
@@ -295,17 +295,17 @@ SELECT create_reference_table('reference_table');
 
 -- show that colociation of citus local tables are not supported for now
 -- between citus local tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['citus_local_table_2']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'citus_local_table_2');
 ERROR:  citus local tables cannot be colocated with other tables
 -- between citus local tables and reference tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['reference_table']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'reference_table');
 ERROR:  citus local tables cannot be colocated with other tables
-SELECT mark_tables_colocated('reference_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('reference_table', colocate_with => 'citus_local_table_1');
 ERROR:  citus local tables cannot be colocated with other tables
 -- between citus local tables and distributed tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['distributed_table']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'distributed_table');
 ERROR:  citus local tables cannot be colocated with other tables
-SELECT mark_tables_colocated('distributed_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('distributed_table', colocate_with => 'citus_local_table_1');
 ERROR:  citus local tables cannot be colocated with other tables
 -- master_create_empty_shard is not supported
 SELECT master_create_empty_shard('citus_local_table_1');
diff --git a/src/test/regress/expected/multi_colocation_utils.out b/src/test/regress/expected/multi_colocation_utils.out
index cea7003e9..dfc26c7a2 100644
--- a/src/test/regress/expected/multi_colocation_utils.out
+++ b/src/test/regress/expected/multi_colocation_utils.out
@@ -833,7 +833,7 @@ ORDER BY
  table2_groupf                   | 1300081 | t            |    57638 |               |
 (108 rows)
 
--- reset colocation ids to test mark_tables_colocated
+-- reset colocation ids to test update_distributed_table_colocation
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 1;
 DELETE FROM pg_dist_colocation
     WHERE colocationid >= 1 AND colocationid < 1000;
@@ -855,21 +855,20 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
 (0 rows)
 
 -- first check failing cases
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupC']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupc
-DETAIL:  Distribution column types don't match for table1_groupb and table1_groupc.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupD']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupd
-DETAIL:  Shard counts don't match for table1_groupb and table1_groupd.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupE']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupe
-DETAIL:  Shard 1300026 of table1_groupb and shard xxxxx of table1_groupe have different number of shard placements.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupF']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupf
-DETAIL:  Replication models don't match for table1_groupb and table1_groupf.
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB', 'table1_groupD']);
-ERROR:  cannot colocate tables table1_groupb and table1_groupd
-DETAIL:  Shard counts don't match for table1_groupb and table1_groupd.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupC');
+ERROR:  cannot colocate tables table1_groupc and table1_groupb
+DETAIL:  Distribution column types don't match for table1_groupc and table1_groupb.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
+ERROR:  cannot colocate tables table1_groupd and table1_groupb
+DETAIL:  Shard counts don't match for table1_groupd and table1_groupb.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupE');
+ERROR:  cannot colocate tables table1_groupe and table1_groupb
+DETAIL:  Shard 1300058 of table1_groupe and shard xxxxx of table1_groupb have different number of shard placements.
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupF');
+ERROR:  relation table1_groupf should be a hash distributed table
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
+ERROR:  cannot colocate tables table1_groupd and table1_groupb
+DETAIL:  Shard counts don't match for table1_groupd and table1_groupb.
 -- check metadata to see failing calls didn't have any side effects
 SELECT * FROM pg_dist_colocation
     WHERE colocationid >= 1 AND colocationid < 1000
@@ -886,39 +885,39 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
 (0 rows)
 
 -- check successfully cololated tables
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table2_groupB');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupC', ARRAY['table2_groupC']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupC', colocate_with => 'table2_groupC');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupD', ARRAY['table2_groupD']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupD', colocate_with => 'table2_groupD');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupE', ARRAY['table2_groupE', 'table3_groupE']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table2_groupE');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupF', ARRAY['table2_groupF']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table3_groupE');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
 -- check to colocate with itself
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupB']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupB');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
@@ -944,40 +943,53 @@ SELECT * FROM pg_dist_colocation
     ORDER BY colocationid;
  colocationid | shardcount | replicationfactor | distributioncolumntype | distributioncolumncollation
 ---------------------------------------------------------------------
-            2 |          2 |                 1 |                     23 |                           0
-            3 |          2 |                 2 |                     25 |                         100
-            4 |          8 |                 2 |                     23 |                           0
+            1 |          2 |                 1 |                     23 |                           0
+            2 |          2 |                 2 |                     25 |                         100
+            3 |          8 |                 2 |                     23 |                           0
+            4 |          2 |                 2 |                     23 |                           0
             5 |          2 |                 2 |                     23 |                           0
-(4 rows)
+(5 rows)
 
 SELECT logicalrelid, colocationid FROM pg_dist_partition
     WHERE colocationid >= 1 AND colocationid < 1000
     ORDER BY colocationid, logicalrelid;
    logicalrelid    | colocationid
 ---------------------------------------------------------------------
- table1_groupb     |            2
- table2_groupb     |            2
- table1_groupc     |            3
- table2_groupc     |            3
- table1_groupd     |            4
- table2_groupd     |            4
+ table1_groupb     |            1
+ table2_groupb     |            1
+ table1_groupc     |            2
+ table2_groupc     |            2
+ table1_groupd     |            3
+ table2_groupd     |            3
+ table2_groupe     |            4
  table1_groupe     |            5
- table2_groupe     |            5
  table3_groupe     |            5
  table1_group_none |            6
  table2_group_none |            7
 (11 rows)
 
 -- move the all tables in colocation group 5 to colocation group 7
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table1_groupE', 'table2_groupE', 'table3_groupE']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table1_groupE');
+ update_distributed_table_colocation
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_groupE');
+ update_distributed_table_colocation
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table3_groupE');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
 
 -- move a table with a colocation id which is already not in pg_dist_colocation
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table2_group_none']);
- mark_tables_colocated
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_group_none');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
@@ -988,27 +1000,29 @@ SELECT * FROM pg_dist_colocation
     ORDER BY colocationid;
  colocationid | shardcount | replicationfactor | distributioncolumntype | distributioncolumncollation
 ---------------------------------------------------------------------
-            2 |          2 |                 1 |                     23 |                           0
-            3 |          2 |                 2 |                     25 |                         100
-            4 |          8 |                 2 |                     23 |                           0
-(3 rows)
+            1 |          2 |                 1 |                     23 |                           0
+            2 |          2 |                 2 |                     25 |                         100
+            3 |          8 |                 2 |                     23 |                           0
+            4 |          2 |                 2 |                     23 |                           0
+            5 |          2 |                 2 |                     23 |                           0
+(5 rows)
 
 SELECT logicalrelid, colocationid FROM pg_dist_partition
     WHERE colocationid >= 1 AND colocationid < 1000
     ORDER BY colocationid, logicalrelid;
    logicalrelid    | colocationid
 ---------------------------------------------------------------------
- table1_groupb     |            2
- table2_groupb     |            2
- table1_groupc     |            3
- table2_groupc     |            3
- table1_groupd     |            4
- table2_groupd     |            4
- table1_groupe     |            6
- table2_groupe     |            6
- table3_groupe     |            6
- table1_group_none |            6
- table2_group_none |            6
+ table1_groupb     |            1
+ table2_groupb     |            1
+ table1_groupc     |            2
+ table2_groupc     |            2
+ table1_groupd     |            3
+ table2_groupd     |            3
+ table2_groupe     |            4
+ table1_groupe     |            5
+ table3_groupe     |            5
+ table1_group_none |            7
+ table2_group_none |            7
 (11 rows)
 
 -- try to colocate different replication models
@@ -1033,9 +1047,9 @@ SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'NONE');
 
 (1 row)
 
-SELECT mark_tables_colocated('table1_groupG', ARRAY['table2_groupG']);
-ERROR:  cannot colocate tables table1_groupg and table2_groupg
-DETAIL:  Replication models don't match for table1_groupg and table2_groupg.
+SELECT update_distributed_table_colocation('table1_groupG', colocate_with => 'table2_groupG');
+ERROR:  cannot colocate tables table2_groupg and table1_groupg
+DETAIL:  Replication models don't match for table2_groupg and table1_groupg.
 CREATE TABLE d1(a int, b int);
 CREATE TABLE d2(a int, b int);
 CREATE TABLE d3(a int, b int);
@@ -1101,31 +1115,6 @@ SELECT create_reference_table('ref');
 
 (1 row)
 
--- check d1, d2, d3 and d4 has the same colocation id => they are colocated.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1169,31 +1158,6 @@ SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
 (1 row)
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      11
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1226,31 +1190,6 @@ SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
 (1 row)
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1281,31 +1220,6 @@ SELECT update_distributed_table_colocation('d3', colocate_with => 'd2');
 
 (1 row)
 
--- d1 and d4 should be colocated, d2 and d3 should be colocated.
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
@@ -1337,37 +1251,6 @@ SELECT update_distributed_table_colocation('d3', colocate_with => '"none"');
 
 (1 row)
 
--- d1, d4, d3 and "none" should be colocated;
-SELECT get_table_colocation_id('d1');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d2');
- get_table_colocation_id
----------------------------------------------------------------------
-                      12
-(1 row)
-
-SELECT get_table_colocation_id('d3');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('d4');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
-SELECT get_table_colocation_id('none');
- get_table_colocation_id
----------------------------------------------------------------------
-                       8
-(1 row)
-
 SELECT tables_colocated('d1', 'd2');
  tables_colocated
 ---------------------------------------------------------------------
diff --git a/src/test/regress/expected/multi_drop_extension.out b/src/test/regress/expected/multi_drop_extension.out
index cd3098ee8..45ae5e9eb 100644
--- a/src/test/regress/expected/multi_drop_extension.out
+++ b/src/test/regress/expected/multi_drop_extension.out
@@ -16,7 +16,6 @@ SELECT create_distributed_table('testtableddl', 'distributecol', 'append');
 -- around.
 SET client_min_messages TO 'WARNING';
 DROP FUNCTION pg_catalog.master_create_worker_shards(text, integer, integer);
-DROP FUNCTION pg_catalog.mark_tables_colocated(regclass, regclass[]);
 DROP EXTENSION citus CASCADE;
 RESET client_min_messages;
 BEGIN;
diff --git a/src/test/regress/expected/multi_metadata_sync.out b/src/test/regress/expected/multi_metadata_sync.out
index c457b5a6b..8d1452be9 100644
--- a/src/test/regress/expected/multi_metadata_sync.out
+++ b/src/test/regress/expected/multi_metadata_sync.out
@@ -808,13 +808,8 @@ SELECT "Constraint", "Definition" FROM table_fkeys WHERE relid='mx_test_schema_1
  mx_fk_constraint_2 | FOREIGN KEY (col1) REFERENCES mx_test_schema_2.mx_table_2(col1)
 (1 row)
 
--- Check that mark_tables_colocated call propagates the changes to the workers
+-- Check that update_distributed_table_colocation call propagates the changes to the workers
 \c - - - :master_port
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 SELECT nextval('pg_catalog.pg_dist_colocationid_seq') AS last_colocation_id \gset
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 10000;
 SET citus.shard_count TO 7;
@@ -865,9 +860,9 @@ SET
 WHERE
 	logicalrelid = 'mx_colocation_test_1'::regclass
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
--- Mark tables colocated and see the changes on the master and the worker
-SELECT mark_tables_colocated('mx_colocation_test_1', ARRAY['mx_colocation_test_2']);
- mark_tables_colocated
+-- Update colocation and see the changes on the master and the worker
+SELECT update_distributed_table_colocation('mx_colocation_test_1', colocate_with => 'mx_colocation_test_2');
+ update_distributed_table_colocation
 ---------------------------------------------------------------------
 
 (1 row)
@@ -881,8 +876,8 @@ WHERE
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
      logicalrelid     | colocationid
 ---------------------------------------------------------------------
- mx_colocation_test_1 |        10001
  mx_colocation_test_2 |        10001
+ mx_colocation_test_1 |        10001
 (2 rows)
 
 \c - - - :worker_1_port
@@ -895,8 +890,8 @@ WHERE
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
      logicalrelid     | colocationid
 ---------------------------------------------------------------------
- mx_colocation_test_1 |        10001
  mx_colocation_test_2 |        10001
+ mx_colocation_test_1 |        10001
 (2 rows)
 
 \c - - - :master_port
diff --git a/src/test/regress/expected/multi_multiuser.out b/src/test/regress/expected/multi_multiuser.out
index daf86bad7..a0480d8a0 100644
--- a/src/test/regress/expected/multi_multiuser.out
+++ b/src/test/regress/expected/multi_multiuser.out
@@ -313,7 +313,7 @@ ABORT;
 SELECT * FROM citus_stat_statements_reset();
 ERROR:  permission denied for function citus_stat_statements_reset
 -- should not be allowed to co-located tables
-SELECT mark_tables_colocated('test', ARRAY['test_coloc'::regclass]);
+SELECT update_distributed_table_colocation('test', colocate_with => 'test_coloc');
 ERROR:  must be owner of table test
 -- should not be allowed to take any locks
 BEGIN;
diff --git a/src/test/regress/expected/multi_reference_table.out b/src/test/regress/expected/multi_reference_table.out
index 83e7e8c84..22ab25eca 100644
--- a/src/test/regress/expected/multi_reference_table.out
+++ b/src/test/regress/expected/multi_reference_table.out
@@ -1196,23 +1196,12 @@ DEBUG:  cannot perform distributed INSERT INTO ... SELECT because the partition
 DETAIL:  The target table's partition column should correspond to a partition column in the subquery.
 DEBUG:  performing repartitioned INSERT ... SELECT
 RESET client_min_messages;
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
--- some tests for mark_tables_colocated
--- should error out
-SELECT mark_tables_colocated('colocated_table_test_2', ARRAY['reference_table_test']);
-ERROR:  cannot colocate tables colocated_table_test_2 and reference_table_test
-DETAIL:  Replication models don't match for colocated_table_test_2 and reference_table_test.
--- should work sliently
-SELECT mark_tables_colocated('reference_table_test', ARRAY['reference_table_test_fifth']);
- mark_tables_colocated
----------------------------------------------------------------------
-
-(1 row)
-
+-- some tests for mark_tables_colocated	
+-- should error out	
+SELECT update_distributed_table_colocation('colocated_table_test_2', colocate_with => 'reference_table_test');
+ERROR:  relation reference_table_test should be a hash distributed table
+SELECT update_distributed_table_colocation('reference_table_test', colocate_with => 'reference_table_test_fifth');
+ERROR:  relation reference_table_test_fifth should be a hash distributed table
 -- ensure that reference tables on
 -- different queries works as expected
 CREATE SCHEMA reference_schema;
diff --git a/src/test/regress/expected/multi_unsupported_worker_operations.out b/src/test/regress/expected/multi_unsupported_worker_operations.out
index 7ec8728bd..ceb34e252 100644
--- a/src/test/regress/expected/multi_unsupported_worker_operations.out
+++ b/src/test/regress/expected/multi_unsupported_worker_operations.out
@@ -220,14 +220,8 @@ SELECT master_remove_node('localhost', 5432);
 (1 row)
 
 \c - - - :worker_1_port
--- mark_tables_colocated
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 UPDATE pg_dist_partition SET colocationid = 0 WHERE logicalrelid='mx_table_2'::regclass;
-SELECT mark_tables_colocated('mx_table', ARRAY['mx_table_2']);
+SELECT update_distributed_table_colocation('mx_table', colocate_with => 'mx_table_2');
 ERROR:  operation is not allowed on this node
 HINT:  Connect to the coordinator and run it again.
 SELECT colocationid FROM pg_dist_partition WHERE logicalrelid='mx_table_2'::regclass;
diff --git a/src/test/regress/sql/citus_local_tables.sql b/src/test/regress/sql/citus_local_tables.sql
index b4e7a7c26..e01534f34 100644
--- a/src/test/regress/sql/citus_local_tables.sql
+++ b/src/test/regress/sql/citus_local_tables.sql
@@ -220,15 +220,15 @@ SELECT create_reference_table('reference_table');
 
 -- show that colociation of citus local tables are not supported for now
 -- between citus local tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['citus_local_table_2']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'citus_local_table_2');
 
 -- between citus local tables and reference tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['reference_table']);
-SELECT mark_tables_colocated('reference_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'reference_table');
+SELECT update_distributed_table_colocation('reference_table', colocate_with => 'citus_local_table_1');
 
 -- between citus local tables and distributed tables
-SELECT mark_tables_colocated('citus_local_table_1', ARRAY['distributed_table']);
-SELECT mark_tables_colocated('distributed_table', ARRAY['citus_local_table_1']);
+SELECT update_distributed_table_colocation('citus_local_table_1', colocate_with => 'distributed_table');
+SELECT update_distributed_table_colocation('distributed_table', colocate_with => 'citus_local_table_1');
 
 -- master_create_empty_shard is not supported
 SELECT master_create_empty_shard('citus_local_table_1');
diff --git a/src/test/regress/sql/multi_colocation_utils.sql b/src/test/regress/sql/multi_colocation_utils.sql
index 55f563234..4959087f8 100644
--- a/src/test/regress/sql/multi_colocation_utils.sql
+++ b/src/test/regress/sql/multi_colocation_utils.sql
@@ -344,7 +344,7 @@ ORDER BY
     shardid,
     nodeport;
 
--- reset colocation ids to test mark_tables_colocated
+-- reset colocation ids to test update_distributed_table_colocation
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 1;
 DELETE FROM pg_dist_colocation
     WHERE colocationid >= 1 AND colocationid < 1000;
@@ -361,11 +361,11 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
     ORDER BY colocationid, logicalrelid;
 
 -- first check failing cases
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupC']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupD']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupE']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupF']);
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB', 'table1_groupD']);
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupC');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupE');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupF');
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupD');
 
 -- check metadata to see failing calls didn't have any side effects
 SELECT * FROM pg_dist_colocation
@@ -377,14 +377,14 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
     ORDER BY colocationid, logicalrelid;
 
 -- check successfully cololated tables
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB']);
-SELECT mark_tables_colocated('table1_groupC', ARRAY['table2_groupC']);
-SELECT mark_tables_colocated('table1_groupD', ARRAY['table2_groupD']);
-SELECT mark_tables_colocated('table1_groupE', ARRAY['table2_groupE', 'table3_groupE']);
-SELECT mark_tables_colocated('table1_groupF', ARRAY['table2_groupF']);
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table2_groupB');
+SELECT update_distributed_table_colocation('table1_groupC', colocate_with => 'table2_groupC');
+SELECT update_distributed_table_colocation('table1_groupD', colocate_with => 'table2_groupD');
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table2_groupE');
+SELECT update_distributed_table_colocation('table1_groupE', colocate_with => 'table3_groupE');
 
 -- check to colocate with itself
-SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupB']);
+SELECT update_distributed_table_colocation('table1_groupB', colocate_with => 'table1_groupB');
 
 SET citus.shard_count = 2;
 
@@ -404,10 +404,12 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
     ORDER BY colocationid, logicalrelid;
 
 -- move the all tables in colocation group 5 to colocation group 7
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table1_groupE', 'table2_groupE', 'table3_groupE']);
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table1_groupE');
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_groupE');
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table3_groupE');
 
 -- move a table with a colocation id which is already not in pg_dist_colocation
-SELECT mark_tables_colocated('table1_group_none', ARRAY['table2_group_none']);
+SELECT update_distributed_table_colocation('table1_group_none', colocate_with => 'table2_group_none');
 
 -- check metadata to see that unused colocation group is deleted
 SELECT * FROM pg_dist_colocation
@@ -431,7 +433,7 @@ SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'table1_
 CREATE TABLE table2_groupG ( id int );
 SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'NONE');
 
-SELECT mark_tables_colocated('table1_groupG', ARRAY['table2_groupG']);
+SELECT update_distributed_table_colocation('table1_groupG', colocate_with => 'table2_groupG');
 
 CREATE TABLE d1(a int, b int);
 CREATE TABLE d2(a int, b int);
@@ -457,14 +459,6 @@ SELECT create_distributed_table('range_table', 'a', 'range');
 
 SELECT create_reference_table('ref');
 
-
--- check d1, d2, d3 and d4 has the same colocation id => they are colocated.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d2', 'd4');
@@ -475,12 +469,6 @@ SELECT tables_colocated('d1', 'd4');
 -- break colocation of d2
 SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d1', 'd3');
@@ -490,12 +478,6 @@ SELECT tables_colocated('d1', 'd4');
 -- update colocation should not error if d2 doesn't have any colocated table.
 SELECT update_distributed_table_colocation('d2', colocate_with => 'none');
 
--- d1 and d3 and d4 should be colocated, d2 should have a new colocation id.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d1', 'd3');
@@ -503,12 +485,6 @@ SELECT tables_colocated('d1', 'd4');
 
 SELECT update_distributed_table_colocation('d3', colocate_with => 'd2');
 
--- d1 and d4 should be colocated, d2 and d3 should be colocated.
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
 SELECT tables_colocated('d1', 'd3');
@@ -516,12 +492,6 @@ SELECT tables_colocated('d1', 'd4');
 
 -- special case, colocate with a table named "none".
 SELECT update_distributed_table_colocation('d3', colocate_with => '"none"');
--- d1, d4, d3 and "none" should be colocated;
-SELECT get_table_colocation_id('d1');
-SELECT get_table_colocation_id('d2');
-SELECT get_table_colocation_id('d3');
-SELECT get_table_colocation_id('d4');
-SELECT get_table_colocation_id('none');
 
 SELECT tables_colocated('d1', 'd2');
 SELECT tables_colocated('d2', 'd3');
diff --git a/src/test/regress/sql/multi_drop_extension.sql b/src/test/regress/sql/multi_drop_extension.sql
index ec1bee234..0598d31ed 100644
--- a/src/test/regress/sql/multi_drop_extension.sql
+++ b/src/test/regress/sql/multi_drop_extension.sql
@@ -16,7 +16,6 @@ SELECT create_distributed_table('testtableddl', 'distributecol', 'append');
 -- around.
 SET client_min_messages TO 'WARNING';
 DROP FUNCTION pg_catalog.master_create_worker_shards(text, integer, integer);
-DROP FUNCTION pg_catalog.mark_tables_colocated(regclass, regclass[]);
 DROP EXTENSION citus CASCADE;
 RESET client_min_messages;
 
diff --git a/src/test/regress/sql/multi_metadata_sync.sql b/src/test/regress/sql/multi_metadata_sync.sql
index a0ff53a7c..8a9b380af 100644
--- a/src/test/regress/sql/multi_metadata_sync.sql
+++ b/src/test/regress/sql/multi_metadata_sync.sql
@@ -354,13 +354,8 @@ NOT VALID;
 \c - - - :worker_1_port
 SELECT "Constraint", "Definition" FROM table_fkeys WHERE relid='mx_test_schema_1.mx_table_1'::regclass;
 
--- Check that mark_tables_colocated call propagates the changes to the workers
+-- Check that update_distributed_table_colocation call propagates the changes to the workers
 \c - - - :master_port
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 SELECT nextval('pg_catalog.pg_dist_colocationid_seq') AS last_colocation_id \gset
 ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 10000;
 SET citus.shard_count TO 7;
@@ -400,8 +395,8 @@ WHERE
 	logicalrelid = 'mx_colocation_test_1'::regclass
 	OR logicalrelid = 'mx_colocation_test_2'::regclass;
 
--- Mark tables colocated and see the changes on the master and the worker
-SELECT mark_tables_colocated('mx_colocation_test_1', ARRAY['mx_colocation_test_2']);
+-- Update colocation and see the changes on the master and the worker
+SELECT update_distributed_table_colocation('mx_colocation_test_1', colocate_with => 'mx_colocation_test_2');
 SELECT
 	logicalrelid, colocationid
 FROM
diff --git a/src/test/regress/sql/multi_multiuser.sql b/src/test/regress/sql/multi_multiuser.sql
index 90cd28790..b88fadb03 100644
--- a/src/test/regress/sql/multi_multiuser.sql
+++ b/src/test/regress/sql/multi_multiuser.sql
@@ -193,7 +193,7 @@ ABORT;
 SELECT * FROM citus_stat_statements_reset();
 
 -- should not be allowed to co-located tables
-SELECT mark_tables_colocated('test', ARRAY['test_coloc'::regclass]);
+SELECT update_distributed_table_colocation('test', colocate_with => 'test_coloc');
 
 -- should not be allowed to take any locks
 BEGIN;
diff --git a/src/test/regress/sql/multi_reference_table.sql b/src/test/regress/sql/multi_reference_table.sql
index 34b9900db..a432730d5 100644
--- a/src/test/regress/sql/multi_reference_table.sql
+++ b/src/test/regress/sql/multi_reference_table.sql
@@ -764,18 +764,11 @@ WHERE
 
 RESET client_min_messages;
 
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
-
 -- some tests for mark_tables_colocated
 -- should error out
-SELECT mark_tables_colocated('colocated_table_test_2', ARRAY['reference_table_test']);
+SELECT update_distributed_table_colocation('colocated_table_test_2', colocate_with => 'reference_table_test');
 
--- should work sliently
-SELECT mark_tables_colocated('reference_table_test', ARRAY['reference_table_test_fifth']);
+SELECT update_distributed_table_colocation('reference_table_test', colocate_with => 'reference_table_test_fifth');
 
 -- ensure that reference tables on
 -- different queries works as expected
diff --git a/src/test/regress/sql/multi_unsupported_worker_operations.sql b/src/test/regress/sql/multi_unsupported_worker_operations.sql
index 69677d3c1..5b07f2cac 100644
--- a/src/test/regress/sql/multi_unsupported_worker_operations.sql
+++ b/src/test/regress/sql/multi_unsupported_worker_operations.sql
@@ -142,15 +142,9 @@ SELECT master_remove_node('localhost', 5432);
 
 \c - - - :worker_1_port
 
--- mark_tables_colocated
--- this function is dropped in Citus10, added here for tests
-CREATE OR REPLACE FUNCTION pg_catalog.mark_tables_colocated(source_table_name regclass, target_table_names regclass[])
-	RETURNS void
-	LANGUAGE C STRICT
-	AS 'citus', $$mark_tables_colocated$$;
 UPDATE pg_dist_partition SET colocationid = 0 WHERE logicalrelid='mx_table_2'::regclass;
 
-SELECT mark_tables_colocated('mx_table', ARRAY['mx_table_2']);
+SELECT update_distributed_table_colocation('mx_table', colocate_with => 'mx_table_2');
 SELECT colocationid FROM pg_dist_partition WHERE logicalrelid='mx_table_2'::regclass;
 
 SELECT colocationid AS old_colocation_id